Fix apply

Author: wjsi

mars/dataframe/groupby/aggregation.py

@@ -94,6 +94,7 @@ def get(self):
     "skew": lambda x, bias=False: x.skew(bias=bias),
     "kurt": lambda x, bias=False: x.kurt(bias=bias),
     "kurtosis": lambda x, bias=False: x.kurtosis(bias=bias),
+    "nunique": lambda x: x.nunique(),
 }
 _series_col_name = "col_name"
 
@@ -720,7 +721,8 @@ def _do_custom_agg(op, custom_reduction, *input_objs):
                     result = (result,)
 
             if out.ndim == 2:
-                result = tuple(r.to_frame().T for r in result)
+                if result[0].ndim == 1:
+                    result = tuple(r.to_frame().T for r in result)
                 if op.stage == OperandStage.agg:
                     result = tuple(r.astype(out.dtypes) for r in result)
             else:

mars/dataframe/reduction/aggregation.py

@@ -78,6 +78,7 @@ def where_function(cond, var1, var2):
     "skew": lambda x, skipna=True, bias=False: x.skew(skipna=skipna, bias=bias),
     "kurt": lambda x, skipna=True, bias=False: x.kurt(skipna=skipna, bias=bias),
     "kurtosis": lambda x, skipna=True, bias=False: x.kurtosis(skipna=skipna, bias=bias),
+    "nunique": lambda x: x.nunique(),
 }
 
 

mars/dataframe/reduction/core.py

@@ -972,13 +972,15 @@ def _compile_function(self, func, func_name=None, ndim=1) -> ReductionSteps:
             else:
                 map_func_name, agg_func_name = step_func_name, step_func_name
 
+            op_custom_reduction = getattr(t.op, "custom_reduction", None)
+
             # build agg description
             agg_funcs.append(
                 ReductionAggStep(
                     agg_input_key,
                     map_func_name,
                     agg_func_name,
-                    custom_reduction,
+                    op_custom_reduction or custom_reduction,
                     t.key,
                     output_limit,
                     t.op.get_reduction_args(axis=self._axis),

mars/dataframe/reduction/tests/test_reduction_execution.py

@@ -671,6 +671,12 @@ def test_nunique(setup, check_ref_counts):
     expected = data1.nunique(axis=1)
     pd.testing.assert_series_equal(result, expected)
 
+    # test with agg func
+    df = md.DataFrame(data1, chunk_size=3)
+    result = df.agg("nunique").execute().fetch()
+    expected = data1.agg("nunique")
+    pd.testing.assert_series_equal(result, expected)
+
 
 @pytest.mark.skipif(pa is None, reason="pyarrow not installed")
 def test_use_arrow_dtype_nunique(setup, check_ref_counts):