Fix as_index when calling groupby-agg (#2676)

Author: wjsi

mars/dataframe/groupby/aggregation.py

@@ -281,6 +281,16 @@ def _get_index_levels(self, groupby, mock_index):
             pd_index = mock_index
         return 1 if not isinstance(pd_index, pd.MultiIndex) else len(pd_index.levels)
 
+    def _fix_as_index(self, result_index: pd.Index):
+        # make sure if as_index=False takes effect
+        if isinstance(result_index, pd.MultiIndex):
+            # if MultiIndex, as_index=False definitely takes no effect
+            self.groupby_params["as_index"] = True
+        elif result_index.name is not None:
+            # if not MultiIndex and agg_df.index has a name
+            # means as_index=False takes no effect
+            self.groupby_params["as_index"] = True
+
     def _call_dataframe(self, groupby, input_df):
         agg_df = build_mock_agg_result(
             groupby, self.groupby_params, self.raw_func, **self.raw_func_kw
@@ -291,13 +301,7 @@ def _call_dataframe(self, groupby, input_df):
         index_value.value.should_be_monotonic = True
 
         # make sure if as_index=False takes effect
-        if isinstance(agg_df.index, pd.MultiIndex):
-            # if MultiIndex, as_index=False definitely takes no effect
-            self.groupby_params["as_index"] = True
-        elif agg_df.index.name is not None:
-            # if not MultiIndex and agg_df.index has a name
-            # means as_index=False takes no effect
-            self.groupby_params["as_index"] = True
+        self._fix_as_index(agg_df.index)
 
         # determine num of indices to group in intermediate steps
         self._index_levels = self._get_index_levels(groupby, agg_df.index)
@@ -315,6 +319,10 @@ def _call_series(self, groupby, in_series):
         agg_result = build_mock_agg_result(
             groupby, self.groupby_params, self.raw_func, **self.raw_func_kw
         )
+
+        # make sure if as_index=False takes effect
+        self._fix_as_index(agg_result.index)
+
         index_value = parse_index(
             agg_result.index, groupby.key, groupby.index_value.key
         )

mars/dataframe/groupby/tests/test_groupby_execution.py

@@ -315,6 +315,13 @@ def test_groupby_getitem(setup):
         expected.sort_values(["c1", "c2"]).reset_index(drop=True),
     )
 
+    r = mdf.groupby(["c1", "c2"], as_index=False)["c3"].agg(["sum"])
+    expected = raw.groupby(["c1", "c2"], as_index=False)["c3"].agg(["sum"])
+    pd.testing.assert_frame_equal(
+        r.execute().fetch().sort_values(["c1", "c2"]),
+        expected.sort_values(["c1", "c2"]),
+    )
+
 
 def test_dataframe_groupby_agg(setup):
     agg_funs = [

mars/dataframe/reduction/core.py

@@ -860,12 +860,12 @@ def add_function(self, func, ndim, cols=None, func_name=None):
             self._output_key_to_post_steps[step.output_key] = step
             self._update_col_dict(self._output_key_to_post_cols, step.output_key, cols)
 
-    @functools.lru_cache(100)
-    def _compile_expr_function(self, py_src):
+    def _compile_expr_function(self, py_src: str, local_consts: dict):
         from ... import tensor, dataframe
 
         result_store = dict()
-        global_vars = globals()
+        global_vars = globals().copy()
+        global_vars.update(local_consts)
         global_vars.update(dict(mt=tensor, md=dataframe, array=np.array, nan=np.nan))
         exec(
             py_src, global_vars, result_store
@@ -989,24 +989,24 @@ def _compile_function(self, func, func_name=None, ndim=1) -> ReductionSteps:
                 assert len(initial_inputs) == 1
                 input_key = initial_inputs[0].key
 
-                func_str, _ = self._generate_function_str(t.inputs[0])
+                func_str, _, local_consts = self._generate_function_str(t.inputs[0])
                 pre_funcs.append(
                     ReductionPreStep(
                         input_key,
                         agg_input_key,
                         None,
-                        self._compile_expr_function(func_str),
+                        self._compile_expr_function(func_str, local_consts),
                     )
                 )
         # collect function output after agg
-        func_str, input_keys = self._generate_function_str(func_ret)
+        func_str, input_keys, local_consts = self._generate_function_str(func_ret)
         post_funcs.append(
             ReductionPostStep(
                 input_keys,
                 func_ret.key,
                 func_name,
                 None,
-                self._compile_expr_function(func_str),
+                self._compile_expr_function(func_str, local_consts),
             )
         )
         if len(_func_compile_cache) > 100:  # pragma: no cover
@@ -1034,6 +1034,7 @@ def _generate_function_str(self, out_tileable):
 
         input_key_to_var = OrderedDict()
         local_key_to_var = dict()
+        local_consts_to_val = dict()
         ref_counts = dict()
         ref_visited = set()
         local_lines = []
@@ -1086,7 +1087,12 @@ def _interpret_var(v):
                     # get representation for variables
                     if hasattr(v, "key"):
                         return keys_to_vars[v.key]
-                    return v
+                    elif isinstance(v, (int, bool, str, bytes, np.integer, np.bool_)):
+                        return repr(v)
+                    else:
+                        const_name = f"_const_{len(local_consts_to_val)}"
+                        local_consts_to_val[const_name] = v
+                        return const_name
 
                 func_name = func_name_raw = getattr(t.op, "_func_name", None)
                 rfunc_name = getattr(t.op, "_rfunc_name", func_name)
@@ -1187,6 +1193,7 @@ def _interpret_var(v):
             f"    {lines_str}\n"
             f"    return {local_key_to_var[out_tileable.key]}",
             list(input_key_to_var.keys()),
+            local_consts_to_val,
         )
 
     def compile(self) -> ReductionSteps: