Repository: mars-project/mars

Issue:
Support `to_pandas` in a chunked manner
<!--
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, `df.to_pandas()` will load all the data in the web, and return the data together back to client, this may cause OOM for web, hence we should support it in a chunked manner.


Discussion:

Resolving patch:
diff --git a/mars/core.py b/mars/core.py
--- a/mars/core.py
+++ b/mars/core.py
@@ -402,11 +402,14 @@ def _check_session(self, session, action):
             raise ValueError(
                 f'Tileable object {key} must be executed first before {action}')
 
-    def fetch(self, session=None, **kw):
+    def _fetch(self, session=None, **kw):
         session = self._get_session(session)
         self._check_session(session, 'fetch')
         return session.fetch(self, **kw)
 
+    def fetch(self, session=None, **kw):
+        return self._fetch(session=session, **kw)
+
     def fetch_log(self, session=None, offsets=None, sizes=None):
         session = self._get_session(session)
         self._check_session(session, 'fetch_log')
@@ -423,21 +426,21 @@ class _ExecuteAndFetchMixin:
     def _execute_and_fetch(self, session=None, **kw):
         if session is None and len(self._executed_sessions) > 0:
             session = self._executed_sessions[-1]
-        try:
-            # fetch first, to reduce the potential cost of submitting a graph
-            return self.fetch(session=session)
-        except ValueError:
-            # not execute before
-            wait = kw.pop('wait', True)
 
-            def run():
-                return self.execute(session=session, **kw).fetch(session=session)
+        wait = kw.pop('wait', True)
 
-            if wait:
-                return run()
-            else:
-                thread_executor = ThreadPoolExecutor(1)
-                return thread_executor.submit(run)
+        def run():
+            fetch_kwargs = kw.pop('fetch_kwargs', dict())
+            if len(self._executed_sessions) == 0:
+                # not executed before
+                self.execute(session=session, **kw)
+            return self.fetch(session=session, **fetch_kwargs)
+
+        if wait:
+            return run()
+        else:
+            thread_executor = ThreadPoolExecutor(1)
+            return thread_executor.submit(run)
 
 
 class _ToObjectMixin(_ExecuteAndFetchMixin):
diff --git a/mars/dataframe/core.py b/mars/dataframe/core.py
--- a/mars/dataframe/core.py
+++ b/mars/dataframe/core.py
@@ -378,6 +378,36 @@ def to_pandas(self, session=None, **kw):
         return self._execute_and_fetch(session=session, **kw)
 
 
+class _BatchedFetcher:
+    __slots__ = ()
+
+    def _iter(self, batch_size=1000, session=None):
+        from .indexing.iloc import iloc
+
+        size = self.shape[0]
+        n_batch = ceildiv(size, batch_size)
+
+        if n_batch > 1:
+            for i in range(n_batch):
+                batch_data = iloc(self)[batch_size * i: batch_size * (i + 1)]
+                yield batch_data._fetch(session=session)
+        else:
+            yield self._fetch(session=session)
+
+    def iterbatch(self, batch_size=1000, session=None, **kw):
+        # trigger execution
+        self.execute(session=session, **kw)
+        return self._iter(batch_size=batch_size, session=session)
+
+    def fetch(self, session=None, **kw):
+        batch_size = kw.pop('batch_size', 1000)
+        if len(kw) > 0:  # pragma: no cover
+            raise TypeError(
+                f"'{next(iter(kw))}' is an invalid keyword argument for this function")
+        batches = list(self._iter(batch_size=batch_size, session=session))
+        return pd.concat(batches) if len(batches) > 1 else batches[0]
+
+
 class IndexData(HasShapeTileableData, _ToPandasMixin):
     __slots__ = ()
 
@@ -758,7 +788,7 @@ def from_tensor(in_tensor, index=None, name=None):
         return series_from_tensor(in_tensor, index=index, name=name)
 
 
-class SeriesData(BaseSeriesData):
+class SeriesData(_BatchedFetcher, BaseSeriesData):
     _type_name = 'Series'
 
     @classmethod
@@ -1043,7 +1073,7 @@ def axes(self):
         return [self.index, self.columns]
 
 
-class DataFrameData(BaseDataFrameData):
+class DataFrameData(_BatchedFetcher, BaseDataFrameData):
     _type_name = 'DataFrame'
 
     def _to_str(self, representation=False):
@@ -1114,26 +1144,13 @@ def cls(cls, provider):
             return DataFrameDef
         return super().cls(provider)
 
-    def _iter_wrap(self, method, batch_size, session, **kw):
-        from .indexing.iloc import iloc
-
-        # trigger execution
-        self.execute(session=session)
-
-        size = self.shape[0]
-        n_batch = ceildiv(size, batch_size)
-
-        for i in range(n_batch):
-            batch_data = iloc(self)[size * i: size * (i + 1)] \
-                .fetch(session=session)
-            yield from getattr(batch_data, method)(**kw)
-
     def iterrows(self, batch_size=1000, session=None):
-        return self._iter_wrap('iterrows', batch_size, session)
+        for batch_data in self.iterbatch(batch_size=batch_size, session=session):
+            yield from getattr(batch_data, 'iterrows')()
 
     def itertuples(self, index=True, name='Pandas', batch_size=1000, session=None):
-        return self._iter_wrap('itertuples', batch_size, session,
-                               index=index, name=name)
+        for batch_data in self.iterbatch(batch_size=batch_size, session=session):
+            yield from getattr(batch_data, 'itertuples')(index=index, name=name)
 
 
 class DataFrame(HasShapeTileableEnity, _ToPandasMixin):
diff --git a/mars/dataframe/utils.py b/mars/dataframe/utils.py
--- a/mars/dataframe/utils.py
+++ b/mars/dataframe/utils.py
@@ -896,7 +896,7 @@ def fetch_corner_data(df_or_series, session=None) -> pd.DataFrame:
         index_size = min_rows // 2 + 1
 
     if index_size is None:
-        return df_or_series.fetch(session=session)
+        return df_or_series._fetch(session=session)
     else:
         head = iloc(df_or_series)[:index_size]
         tail = iloc(df_or_series)[-index_size:]


Test patch:
diff --git a/mars/tests/test_session.py b/mars/tests/test_session.py
--- a/mars/tests/test_session.py
+++ b/mars/tests/test_session.py
@@ -224,6 +224,8 @@ def testDataFrameExecuteNotFetch(self):
             sess.fetch(df1)
 
         self.assertIs(df1.execute(), df1)
+        self.assertEqual(len(df1[df1['a'] > 1].to_pandas(fetch_kwargs={'batch_size': 2})), 0)
+        self.assertEqual(len(df1[df1['a'] > 1]['a'].to_pandas(fetch_kwargs={'batch_size': 2})), 0)
 
         # modify result
         executor = sess._sess._executor
@@ -233,6 +235,7 @@ def testDataFrameExecuteNotFetch(self):
         expected.iloc[:2, :2] = data1.iloc[:2, :2] * 3
 
         pd.testing.assert_frame_equal(df1.to_pandas(), expected)
+        pd.testing.assert_frame_equal(df1.to_pandas(fetch_kwargs={'batch_size': 2}), expected)
 
     def testClosedSession(self):
         session = new_session()
@@ -619,6 +622,9 @@ def testDataFrameIter(self):
         raw_data = pd.DataFrame(np.random.randint(1000, size=(20, 10)))
         df = md.DataFrame(raw_data, chunk_size=5)
 
+        for i, batch in enumerate(df.iterbatch(batch_size=15)):
+            pd.testing.assert_frame_equal(batch, raw_data.iloc[i * 15: (i + 1) * 15])
+
         i = 0
         for result_row, expect_row in zip(df.iterrows(batch_size=15),
                                           raw_data.iterrows()):