Auto merge small chunks when df.groupby().apply(func) is doing aggregation (#2708)

Author: Xuye (Chris) Qin

mars/dataframe/groupby/apply.py

@@ -17,11 +17,18 @@
 
 from ... import opcodes
 from ...core import OutputType
+from ...core.context import get_context
 from ...core.custom_log import redirect_custom_log
-from ...serialization.serializables import TupleField, DictField, FunctionField
+from ...serialization.serializables import (
+    BoolField,
+    TupleField,
+    DictField,
+    FunctionField,
+)
 from ...utils import enter_current_session, quiet_stdio
 from ..operands import DataFrameOperandMixin, DataFrameOperand
 from ..utils import (
+    auto_merge_chunks,
     build_empty_df,
     build_empty_series,
     parse_index,
@@ -35,26 +42,13 @@ class GroupByApply(DataFrameOperand, DataFrameOperandMixin):
     _op_type_ = opcodes.APPLY
     _op_module_ = "dataframe.groupby"
 
-    _func = FunctionField("func")
-    _args = TupleField("args")
-    _kwds = DictField("kwds")
-
-    def __init__(self, func=None, args=None, kwds=None, output_types=None, **kw):
-        super().__init__(
-            _func=func, _args=args, _kwds=kwds, _output_types=output_types, **kw
-        )
-
-    @property
-    def func(self):
-        return self._func
+    func = FunctionField("func")
+    args = TupleField("args", default_factory=tuple)
+    kwds = DictField("kwds", default_factory=dict)
+    maybe_agg = BoolField("maybe_agg", default=None)
 
-    @property
-    def args(self):
-        return getattr(self, "_args", None) or ()
-
-    @property
-    def kwds(self):
-        return getattr(self, "_kwds", None) or dict()
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
 
     @classmethod
     @redirect_custom_log
@@ -135,7 +129,14 @@ def tile(cls, op):
             kw["nsplits"] = ((np.nan,) * len(chunks), (out_df.shape[1],))
         else:
             kw["nsplits"] = ((np.nan,) * len(chunks),)
-        return new_op.new_tileables([in_groupby], **kw)
+        ret = new_op.new_tileable([in_groupby], **kw)
+        if not op.maybe_agg:
+            return [ret]
+        else:
+            # auto merge small chunks if df.groupby().apply(func)
+            # may be an aggregation operation
+            yield ret.chunks  # trigger execution for chunks
+            return [auto_merge_chunks(get_context(), ret)]
 
     def _infer_df_func_returns(
         self, in_groupby, in_df, dtypes, dtype=None, name=None, index=None
@@ -147,6 +148,12 @@ def _infer_df_func_returns(
                 self.func, *self.args, **self.kwds
             )
 
+            if len(infer_df) <= 2:
+                # we create mock df with 4 rows, 2 groups
+                # if return df has 2 rows, we assume that
+                # it's an aggregation operation
+                self.maybe_agg = True
+
             # todo return proper index when sort=True is implemented
             index_value = parse_index(infer_df.index[:0], in_df.key, self.func)
 

mars/dataframe/groupby/tests/test_groupby.py

@@ -199,7 +199,11 @@ def apply_series(s):
     assert applied.chunks[0].shape == (np.nan,)
     assert applied.chunks[0].dtype == df1.a.dtype
 
-    applied = tile(mdf.groupby("b").apply(lambda df: df.a.sum()))
+    applied = mdf.groupby("b").apply(lambda df: df.a.sum())
+    assert applied.op.maybe_agg is True
+    # force set to pass test
+    applied.op.maybe_agg = None
+    applied = tile(applied)
     assert applied.dtype == df1.a.dtype
     assert applied.shape == (np.nan,)
     assert applied.op._op_type_ == opcodes.APPLY

mars/dataframe/tests/test_utils.py

@@ -15,6 +15,7 @@
 import operator
 from collections import OrderedDict
 from numbers import Integral
+from typing import List, Dict
 
 import numpy as np
 import pandas as pd
@@ -24,7 +25,7 @@
 from ...core import tile
 from ...utils import Timer
 from ..core import IndexValue
-from ..initializer import DataFrame, Index
+from ..initializer import DataFrame, Series, Index
 from ..utils import (
     decide_dataframe_chunk_sizes,
     decide_series_chunk_size,
@@ -39,6 +40,7 @@
     make_dtypes,
     build_concatenated_rows_frame,
     merge_index_value,
+    auto_merge_chunks,
 )
 
 
@@ -582,3 +584,56 @@ def test_build_concatenated_rows_frame(setup, columns):
             concatenated.chunks[i].columns_value.to_pandas(), df.columns
         )
     pd.testing.assert_frame_equal(concatenated.execute().fetch(), df)
+
+
+def test_auto_merge_chunks():
+    from ..merge import DataFrameConcat
+
+    pdf = pd.DataFrame(np.random.rand(16, 4), columns=list("abcd"))
+    memory_size = pdf.iloc[:4].memory_usage().sum()
+
+    class FakeContext:
+        def __init__(self, retval=True):
+            self._retval = retval
+
+        def get_chunks_meta(self, data_keys: List[str], **_) -> List[Dict]:
+            if self._retval:
+                return [{"memory_size": memory_size}] * len(data_keys)
+            else:
+                return [None] * len(data_keys)
+
+    df = tile(DataFrame(pdf, chunk_size=4))
+    df2 = auto_merge_chunks(FakeContext(), df, 2 * memory_size)
+    assert len(df2.chunks) == 2
+    assert isinstance(df2.chunks[0].op, DataFrameConcat)
+    assert len(df2.chunks[0].op.inputs) == 2
+    assert isinstance(df2.chunks[1].op, DataFrameConcat)
+    assert len(df2.chunks[1].op.inputs) == 2
+
+    df2 = auto_merge_chunks(FakeContext(), df, 3 * memory_size)
+    assert len(df2.chunks) == 2
+    assert isinstance(df2.chunks[0].op, DataFrameConcat)
+    assert len(df2.chunks[0].op.inputs) == 3
+    assert df2.chunks[1] is df.chunks[-1]
+
+    # mock situation that df not executed
+    df2 = auto_merge_chunks(FakeContext(False), df, 3 * memory_size)
+    assert df2 is df
+
+    # number of chunks on columns > 1
+    df3 = tile(DataFrame(pdf, chunk_size=2))
+    df4 = auto_merge_chunks(FakeContext(), df3, 2 * memory_size)
+    assert df4 is df3
+
+    # test series
+    ps = pdf.loc[:, "a"]
+    memory_size = ps.iloc[:4].memory_usage()
+    s = tile(Series(ps, chunk_size=4))
+    s2 = auto_merge_chunks(FakeContext(), s, 2 * memory_size)
+    assert len(s2.chunks) == 2
+    assert isinstance(s2.chunks[0].op, DataFrameConcat)
+    assert s2.chunks[0].name == "a"
+    assert len(s2.chunks[0].op.inputs) == 2
+    assert isinstance(s2.chunks[1].op, DataFrameConcat)
+    assert s2.chunks[1].name == "a"
+    assert len(s2.chunks[1].op.inputs) == 2

mars/dataframe/utils.py

@@ -17,17 +17,28 @@
 import operator
 from contextlib import contextmanager
 from numbers import Integral
+from typing import List, Union
 
 import numpy as np
 import pandas as pd
 from pandas.api.types import is_string_dtype
 from pandas.api.extensions import ExtensionDtype
 from pandas.core.dtypes.cast import find_common_type
 
+from ..config import options
 from ..core import Entity, ExecutableTuple
+from ..core.context import Context
 from ..lib.mmh3 import hash as mmh_hash
 from ..tensor.utils import dictify_chunk_size, normalize_chunk_sizes
-from ..utils import tokenize, sbytes, lazy_import, ModulePlaceholder, is_full_slice
+from ..typing import ChunkType, TileableType
+from ..utils import (
+    tokenize,
+    sbytes,
+    lazy_import,
+    ModulePlaceholder,
+    is_full_slice,
+    parse_readable_size,
+)
 
 try:
     import pyarrow as pa
@@ -1293,3 +1304,91 @@ def is_cudf(x):
         if isinstance(x, (cudf.DataFrame, cudf.Series, cudf.Index)):
             return True
     return False
+
+
+def auto_merge_chunks(
+    ctx: Context,
+    df_or_series: TileableType,
+    merged_file_size: Union[int, float, str] = None,
+) -> TileableType:
+    from .merge import DataFrameConcat
+
+    if df_or_series.ndim == 2 and df_or_series.chunk_shape[1] > 1:
+        # skip auto merge optimization for DataFrame
+        # that has more than 1 chunks on columns axis
+        return df_or_series
+
+    metas = ctx.get_chunks_meta(
+        [c.key for c in df_or_series.chunks], fields=["memory_size"], error="ignore"
+    )
+    memory_sizes = [meta["memory_size"] if meta is not None else None for meta in metas]
+    if any(size is None for size in memory_sizes):
+        # has not been executed before, cannot get accurate memory size, skip auto merge
+        return df_or_series
+
+    def _concat_chunks(merge_chunks: List[ChunkType], output_index: int):
+        chunk_size = sum(c.shape[0] for c in merge_chunks)
+        concat_op = DataFrameConcat(output_types=df_or_series.op.output_types)
+        if df_or_series.ndim == 1:
+            kw = dict(
+                dtype=df_or_series.dtype,
+                index_value=merge_index_value(
+                    {c.index: c.index_value for c in merge_chunks}
+                ),
+                shape=(chunk_size,),
+                index=(output_index,),
+                name=df_or_series.name,
+            )
+        else:
+            kw = dict(
+                dtypes=merge_chunks[0].dtypes,
+                index_value=merge_index_value(
+                    {c.index: c.index_value for c in merge_chunks}
+                ),
+                columns_value=merge_chunks[0].columns_value,
+                shape=(chunk_size, merge_chunks[0].shape[1]),
+                index=(output_index, 0),
+            )
+        return concat_op.new_chunk(merge_chunks, **kw)
+
+    to_merge_size = (
+        parse_readable_size(merged_file_size)[0]
+        if merged_file_size is not None
+        else options.chunk_store_limit
+    )
+    to_merge_chunks = []
+    acc_memory_size = 0
+    n_split = []
+    out_chunks = []
+    for chunk, chunk_memory_size in zip(df_or_series.chunks, memory_sizes):
+        if acc_memory_size + chunk_memory_size > to_merge_size:
+            # adding current chunk would exceed the maximum,
+            # concat previous chunks
+            merged_chunk = _concat_chunks(to_merge_chunks, len(n_split))
+            out_chunks.append(merged_chunk)
+            n_split.append(merged_chunk.shape[0])
+            # reset
+            acc_memory_size = 0
+            to_merge_chunks = []
+
+        to_merge_chunks.append(chunk)
+        acc_memory_size += chunk_memory_size
+    # process the last chunk
+    if len(to_merge_chunks) > 1:
+        merged_chunk = _concat_chunks(to_merge_chunks, len(n_split))
+        out_chunks.append(merged_chunk)
+        n_split.append(merged_chunk.shape[0])
+    else:
+        assert len(to_merge_chunks) == 1
+        last_chunk = to_merge_chunks[0]
+        out_chunks.append(last_chunk)
+        n_split.append(last_chunk.shape[0])
+
+    new_op = df_or_series.op.copy()
+    params = df_or_series.params.copy()
+    params["chunks"] = out_chunks
+    if df_or_series.ndim == 1:
+        params["nsplits"] = (tuple(n_split),)
+    else:
+        params["nsplits"] = (tuple(n_split), df_or_series.nsplits[1])
+    return new_op.new_tileable(df_or_series.op.inputs, kws=[params])