Fix potential empty chunks when creating DataFrame from pandas (#2987)

Author: Xuye (Chris) Qin

mars/dataframe/datasource/dataframe.py

@@ -14,11 +14,14 @@
 
 import itertools
 
+import pandas as pd
+
 from ... import opcodes as OperandDef
 from ...config import options
 from ...core import OutputType
 from ...serialization.serializables import DataFrameField, SeriesField
 from ...tensor.utils import get_chunk_slices
+from ...utils import estimate_pandas_size
 from ..utils import decide_dataframe_chunk_sizes, parse_index, is_cudf
 from ..operands import DataFrameOperand, DataFrameOperandMixin
 
@@ -61,7 +64,10 @@ def tile(cls, op: "DataFrameDataSource"):
         df = op.outputs[0]
         raw_df = op.data
 
-        memory_usage = raw_df.memory_usage(index=False, deep=True)
+        # estimate column memory usage instead of calling df.memory_usage(deep=True)
+        memory_usage = pd.Series(
+            {c: estimate_pandas_size(s) for c, s in raw_df.iteritems()}
+        )
         chunk_size = df.extra_params.raw_chunk_size or options.chunk_size
         chunk_size = decide_dataframe_chunk_sizes(df.shape, chunk_size, memory_usage)
         chunk_size_idxes = (range(len(size)) for size in chunk_size)

mars/dataframe/datasource/tests/test_datasource.py

@@ -15,13 +15,15 @@
 import os
 import tempfile
 import shutil
+import string
 from collections import OrderedDict
 
 import numpy as np
 import pandas as pd
 import pytest
 
 from .... import tensor as mt
+from ....config import option_context
 from ....core import tile
 from ....tests.core import require_ray
 from ....utils import lazy_import
@@ -138,6 +140,20 @@ def test_from_pandas_dataframe():
     pd.testing.assert_frame_equal(df2.chunks[5].op.data, df2.op.data.iloc[4:, 8:])
     assert df2.chunks[3].index_value._index_value._slice == slice(8, 10, 2)
 
+    raw = pd.DataFrame(
+        {
+            "a": [
+                string.printable[i : i + 15]
+                for i in np.random.randint(len(string.printable), size=100)
+            ],
+            "b": np.random.rand(100),
+        }
+    )
+    with option_context({"chunk_store_limit": raw["a"].memory_usage(deep=True) / 10}):
+        df = tile(from_pandas_df(raw))
+        # see GH#2985, empty chunks are wrongly generated
+        assert len([ns for ns in df.nsplits[1] if ns == 0]) == 0
+
 
 def test_from_pandas_series():
     data = pd.Series(np.random.rand(10), name="a")

mars/dataframe/utils.py

@@ -201,10 +201,10 @@ def decide_dataframe_chunk_sizes(shape, chunk_size, memory_usage):
             int(np.power(max_chunk_size / nbytes_occupied, 1 / float(nleft))), 1
         )
 
-        if col_left_size == 0:
+        if col_left_size == 0 and not col_chunk_size:
             col_chunk_size.append(0)
 
-        if row_left_size == 0:
+        if row_left_size == 0 and not row_chunk_size:
             row_chunk_size.append(0)
 
         # check col first