Fix reduction agg with UDF (#737)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
xorbitsai · Oct 10, 2023 · 7ea4316 · 7ea4316
1 parent 30cac74
commit 7ea4316
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/python/xorbits/_mars/dataframe/reduction/sum.py b/python/xorbits/_mars/dataframe/reduction/sum.py
@@ -46,7 +46,14 @@ def sum_(value):
 
 
 def sum_series(
- df, axis=None, skipna=True, level=None, min_count=0, combine_size=None, method=None
+ df,
+ axis=None,
+ skipna=True,
+ level=None,
+ min_count=0,
+ combine_size=None,
+ method=None,
+ **kwargs, # kwargs for compatible with numpy reduction
 ):
  op = DataFrameSum(
  axis=axis,
@@ -69,6 +76,7 @@ def sum_dataframe(
  numeric_only=None,
  combine_size=None,
  method=None,
+ **kwargs, # kwargs for compatible with numpy reduction
 ):
  op = DataFrameSum(
  axis=axis,

diff --git a/python/xorbits/_mars/dataframe/reduction/tests/test_reduction_execution.py b/python/xorbits/_mars/dataframe/reduction/tests/test_reduction_execution.py
@@ -771,6 +771,16 @@ def test_dataframe_aggregate(setup, check_ref_counts):
  ]
  data = pd.DataFrame(np.random.rand(20, 20))
 
+ def realized_volatility(series):
+ print(series)
+ return np.sqrt(np.sum(series**2))
+
+ df = md.DataFrame(data)
+ result = df.agg(realized_volatility)
+ pd.testing.assert_series_equal(
+ result.execute().fetch(), data.agg(realized_volatility)
+ )
+
  df = md.DataFrame(data)
  result = df.agg(all_aggs)
  pd.testing.assert_frame_equal(result.execute().fetch(), data.agg(all_aggs))