ENH: len for groupby

python/xorbits/_mars/dataframe/groupby/__init__.py

@@ -12,9 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 # noinspection PyUnresolvedReferences
 from ..core import DataFrameGroupBy, GroupBy, SeriesGroupBy
+from .len import groupby_len
 
 
 def _install():
@@ -63,6 +63,7 @@ def _install():
         setattr(cls, "sem", lambda groupby, **kw: agg(groupby, "sem", **kw))
         setattr(cls, "nunique", lambda groupby, **kw: agg(groupby, "nunique", **kw))
 
+        setattr(cls, "__len__", groupby_len)
         setattr(cls, "apply", groupby_apply)
         setattr(cls, "transform", groupby_transform)
 

python/xorbits/_mars/dataframe/groupby/len.py

@@ -0,0 +1,84 @@
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...core.operand import Operand, OperandStage
+from ..operands import DataFrameOperandMixin
+
+
+class GroupByLen(DataFrameOperandMixin, Operand):
+    _op_type_ = opcodes.GROUPBY_LEN
+
+    def __call__(self, groupby):
+        return self.new_scalar([groupby])
+
+    @classmethod
+    def tile(cls, op: "GroupByLen"):
+        in_groupby = op.inputs[0]
+
+        # generate map chunks
+        map_chunks = []
+        for chunk in in_groupby.chunks:
+            map_op = op.copy().reset_key()
+            map_op.stage = OperandStage.map
+            map_op.output_types = [OutputType.series]
+            chunk_inputs = [chunk]
+
+            map_chunks.append(map_op.new_chunk(chunk_inputs))
+
+        # generate reduce chunks, we only need one reducer here.
+        out_chunks = []
+        reduce_op = op.copy().reset_key()
+        reduce_op.output_types = [OutputType.scalar]
+        reduce_op.stage = OperandStage.reduce
+        out_chunks.append(reduce_op.new_chunk(map_chunks))
+
+        # final wrap up:
+        new_op = op.copy()
+        params = op.outputs[0].params.copy()
+        params["nsplits"] = ((np.nan,) * len(out_chunks),)
+        params["chunks"] = out_chunks
+        return new_op.new_scalar(op.inputs, **params)
+
+    @classmethod
+    def execute_map(cls, ctx, op):
+        chunk = op.outputs[0]
+        in_df_grouped = ctx[op.inputs[0].key]
+
+        # grouped object .size() method ensure every unique keys
+        summary = in_df_grouped.size()
+        sum_indexes = summary.index
+
+        res = []
+        for index in sum_indexes:
+            res.append(index)
+
+        # use series to convey every index store in this level
+        ctx[chunk.key, 1] = pd.Series(res)
+
+    @classmethod
+    def execute_reduce(cls, ctx, op: "GroupByLen"):
+        chunk = op.outputs[0]
+        input_idx_to_series = dict(op.iter_mapper_data(ctx))
+        row_idxes = sorted(input_idx_to_series.keys())
+
+        res = set()
+        for row_index in row_idxes:
+            row_series = input_idx_to_series.get(row_index, None)
+            res.update(row_series)
+
+        res_len = len(res)
+        ctx[chunk.key] = res_len
+
+    @classmethod
+    def execute(cls, ctx, op: "GroupByLen"):
+        if op.stage == OperandStage.map:
+            cls.execute_map(ctx, op)
+        elif op.stage == OperandStage.reduce:
+            cls.execute_reduce(ctx, op)
+
+
+def groupby_len(groupby):
+    op = GroupByLen()
+    return op(groupby).execute().fetch()

python/xorbits/_mars/dataframe/groupby/tests/test_groupby.py

@@ -525,3 +525,18 @@ def test_groupby_fill():
     assert len(r.chunks) == 4
     assert r.shape == (len(s1),)
     assert r.chunks[0].shape == (np.nan,)
+
+
+def test_groupby_len(setup):
+    df = md.DataFrame(
+        {
+            "a": ["a", "b", "a", "c"],
+            "b": [0.1, 0.2, 0.3, 0.4],
+            "c": ["aa", "bb", "cc", "aa"],
+        }
+    )
+
+    grouped = df.groupby("b")
+
+    num_groups = len(grouped)
+    print(num_groups)

python/xorbits/_mars/dataframe/groupby/tests/test_groupby_execution.py

@@ -1886,3 +1886,16 @@ def test_series_groupby_rolling_agg(setup, window, min_periods, center, closed,
     mresult = mresult.execute().fetch()
 
     pd.testing.assert_series_equal(presult, mresult.sort_index())
+
+
+def test_grouby_len(setup):
+    df = md.DataFrame(
+        {
+            "a": ["a", "b", "a", "c"],
+            "b": [0.1, 0.2, 0.3, 0.4],
+            "c": ["aa", "bb", "cc", "aa"],
+        }
+    )
+    grouped = df.groupby("b")
+
+    print(len(grouped))

python/xorbits/_mars/opcodes.py

@@ -390,6 +390,7 @@
 APPLYMAP = 742
 PIVOT = 743
 PIVOT_TABLE = 744
+LEN = 745
 
 FUSE = 801
 
@@ -434,6 +435,7 @@
 GROUPBY_SORT_REGULAR_SAMPLE = 2037
 GROUPBY_SORT_PIVOT = 2038
 GROUPBY_SORT_SHUFFLE = 2039
+GROUPBY_LEN = 2064
 
 # parallel sorting by regular sampling
 PSRS_SORT_REGULAR_SMAPLE = 2040