Merge branch 'main' into bug/column_pruning_groupby_multi_series

python/xorbits/_mars/dataframe/reduction/aggregation.py

@@ -952,7 +952,11 @@ def execute(cls, ctx, op: "DataFrameAggregate"):
  if is_cudf(in_data):
  result = cls._cudf_agg(op, in_data)
  else:
- result = in_data.agg(op.raw_func, axis=op.axis)
+ result = (
+ in_data.agg(op.raw_func, axis=op.axis)
+ if op.raw_func is not None
+ else in_data.agg(**op.raw_func_kw, axis=op.axis)
+ )
  if op.outputs[0].ndim == 1:
  result = result.astype(op.outputs[0].dtype, copy=False)
 

python/xorbits/_mars/dataframe/reduction/tests/test_reduction_execution.py

@@ -1138,3 +1138,34 @@ def g3(x):
  s.agg((g1, g2, g3)), ms.agg((g1, g2, g3)).execute().fetch()
  )
  pd.testing.assert_series_equal(s.agg((g1, g1)), ms.agg((g1, g1)).execute().fetch())
+
+
+@pytest.mark.parametrize("chunk_size", [None, 1, 5, 10])
+def test_agg_with_kwargs(setup, chunk_size):
+ rs = np.random.RandomState(0)
+ df = pd.DataFrame(
+ {
+ "a": rs.choice([1, 3, 8], size=100),
+ "b": rs.choice([201.8, 155.7, 95.7], size=100),
+ "c": rs.choice([1, np.nan, 3], size=100),
+ },
+ )
+ mdf = md.DataFrame(df, chunk_size=chunk_size)
+ res = mdf.agg(a=("a", "sum"))
+ pd.testing.assert_frame_equal(res.execute().fetch(), df.agg(a=("a", "sum")))
+
+ res = mdf.agg(x=("a", "sum"), y=("b", "mean"))
+ pd.testing.assert_frame_equal(
+ res.execute().fetch(), df.agg(x=("a", "sum"), y=("b", "mean"))
+ )
+
+ res = mdf.agg(x=("a", "mean"), y=("c", sum))
+ pd.testing.assert_frame_equal(
+ res.execute().fetch(), df.agg(x=("a", "mean"), y=("c", sum))
+ )
+
+ def g(x):
+ return x.sum() - (x * 3).sum()
+
+ res = mdf.agg(g=("b", g))
+ pd.testing.assert_frame_equal(res.execute().fetch(), df.agg(g=("b", g)))