diff --git a/python/xorbits/core/adapter.py b/python/xorbits/core/adapter.py index 96079458d..705aca3d9 100644 --- a/python/xorbits/core/adapter.py +++ b/python/xorbits/core/adapter.py @@ -62,6 +62,7 @@ from .._mars.dataframe.datastore.to_vineyard import ( DataFrameToVineyardChunk as MarsDataFrameToVineyardChunk, ) +from .._mars.dataframe.groupby.rolling import GroupByRolling as MarsGroupByRolling from .._mars.dataframe.indexing.at import DataFrameAt as MarsDataFrameAt from .._mars.dataframe.indexing.iat import DataFrameIat as MarsDataFrameIat from .._mars.dataframe.indexing.iloc import DataFrameIloc as MarsDataFrameIloc diff --git a/python/xorbits/pandas/groupby.py b/python/xorbits/pandas/groupby.py index c28b94b7b..2bd035f71 100644 --- a/python/xorbits/pandas/groupby.py +++ b/python/xorbits/pandas/groupby.py @@ -15,6 +15,7 @@ import pandas from ..core import DataRef, DataType +from ..core.adapter import MarsGetAttrProxy, MarsGroupByRolling, register_converter from ..core.data import register_cls_to_type from ..core.utils.docstring import attach_module_callable_docstring @@ -37,3 +38,8 @@ class SeriesGroupBy(DataRef): attach_module_callable_docstring( SeriesGroupBy, pandas, pandas.core.groupby.SeriesGroupBy ) + + +@register_converter(from_cls_list=[MarsGroupByRolling]) +class GroupByRolling(MarsGetAttrProxy): + pass diff --git a/python/xorbits/pandas/tests/test_groupby.py b/python/xorbits/pandas/tests/test_groupby.py new file mode 100644 index 000000000..f8d5f6e65 --- /dev/null +++ b/python/xorbits/pandas/tests/test_groupby.py @@ -0,0 +1,27 @@ +# Copyright 2022-2023 XProbe Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +import pandas as pd + +from ... import pandas as xpd + + +def test_groupby_rolling(setup): + df = pd.DataFrame({"Group": ["A", "A", "B", "B", "B"], "Value": [1, 2, 3, 4, 5]}) + xdf = xpd.DataFrame(df) + + rolling_x_mean = xdf.groupby("Group")["Value"].rolling(window=2).mean().to_pandas() + rolling_mean = df.groupby("Group")["Value"].rolling(window=2).mean() + + pd.testing.assert_series_equal(rolling_x_mean, rolling_mean)