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,86 @@
+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
+        params = dict(dtype=int)
+        out_chunks.append(
+            reduce_op.new_chunk(map_chunks, shape=(), index=(0,), dtype=int)
+        )
+
+        # final wrap up:
+        new_op = op.copy()
+        params = op.outputs[0].params.copy()
+
+        params.pop("shape")
+
+        params["chunks"] = out_chunks
+        return new_op.new_scalars(op.inputs, **params)
+
+    @classmethod
+    def execute_map(cls, ctx, op: "GroupByLen"):
+        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]
+        key = op.inputs[0].key
+
+        res = set()
+        input_series = ctx[key, 1]
+        res.update(input_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,26 @@ 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):
+    np.random.seed(42)
+    num_dataframes = 10
+    for i in range(num_dataframes):
+        # Generate random data
+        data = {
+            "Category": np.random.choice(["A", "B", "C"], size=100),
+            "Value": np.random.randint(1, 100, size=100),
+        }
+
+        # Create DataFrame
+        df = md.DataFrame(data)
+        df_test = pd.DataFrame(data)
+
+        grouped = df.groupby("Category")
+        grouped_test = df_test.groupby("Category")
+
+        grouped2 = df.groupby("Value")
+        grouped_test2 = df_test.groupby("Value")
+        assert len(grouped) == len(grouped_test)
+        assert len(grouped2) == len(grouped_test2)

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