maxframe 2.0.0b2__cp37-cp37m-win_amd64.whl → 2.2.0__cp37-cp37m-win_amd64.whl

maxframe/dataframe/extensions/map_reduce.py

@@ -0,0 +1,263 @@
+# Copyright 1999-2025 Alibaba Group Holding Ltd.
+#
+# 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 inspect
+from typing import Any, Callable, List, Optional, Union
+
+import pandas as pd
+
+
+def _has_end_arg(func) -> bool:
+    f_args = inspect.getfullargspec(func)
+    return "end" in f_args.args or "end" in f_args.kwonlyargs
+
+
+def _gen_combined_mapper(
+    mapper: Callable,
+    combiner: Callable,
+    group_cols: List[Any],
+    order_cols: List[Any],
+    ascending: Union[bool, List[bool]] = True,
+):
+    class CombinedMapper:
+        def __init__(self):
+            if isinstance(mapper, type):
+                self.f = mapper()
+            else:
+                self.f = mapper
+
+            if isinstance(combiner, type):
+                self.combiner = combiner()
+            else:
+                self.combiner = combiner
+
+        def _combine_mapper_result(self, mapper_result, end=False):
+            if mapper_result is None:
+                return None
+            res = mapper_result
+            if order_cols:
+                res = mapper_result.sort_values(order_cols, ascending=ascending)
+
+            kw = {"end": end} if _has_end_arg(self.combiner) else {}
+            gcols = group_cols or list(res.columns)
+            return res.groupby(gcols, group_keys=False)[list(res.columns)].apply(
+                self.combiner, **kw
+            )
+
+        def __call__(self, batch, end=False):
+            kw = {"end": end} if _has_end_arg(self.f) else {}
+            f_ret = self.f(batch, **kw)
+            return self._combine_mapper_result(f_ret, end=end)
+
+        def close(self) -> None:
+            if hasattr(self.f, "close"):
+                self.f.close()
+            if hasattr(self.combiner, "close"):
+                self.combiner.close()
+
+    return CombinedMapper
+
+
+def map_reduce(
+    df,
+    mapper: Optional[Callable] = None,
+    reducer: Optional[Callable] = None,
+    group_cols: Optional[List[Any]] = None,
+    *,
+    order_cols: List[Any] = None,
+    ascending: Union[bool, List[bool]] = True,
+    combiner: Callable = None,
+    batch_rows: Optional[int] = 1024,
+    mapper_dtypes: pd.Series = None,
+    mapper_index: pd.Index = None,
+    mapper_batch_rows: Optional[int] = None,
+    reducer_dtypes: pd.Series = None,
+    reducer_index: pd.Index = None,
+    reducer_batch_rows: Optional[int] = None,
+    ignore_index: bool = False,
+):
+    """
+    Map-reduce API over certain DataFrames. This function is roughly
+    a shortcut for
+
+    .. code-block:: python
+
+        df.mf.apply_chunk(mapper).groupby(group_keys).mf.apply_chunk(reducer)
+
+    Parameters
+    ----------
+    mapper : function or type
+        Mapper function or class.
+    reducer : function or type
+        Reducer function or class.
+    group_cols : str or list[str]
+        The keys to group after mapper. If absent, all columns in the mapped
+        DataFrame will be used.
+    order_cols : str or list[str]
+        The columns to sort after groupby.
+    ascending : bool or list[bool] or None
+        Whether columns should be in ascending order or not, only effective when
+        `order_cols` are specified. If a list of booleans are passed, orders of
+        every column in `order_cols` are specified.
+    combiner : function or class
+        Combiner function or class. Should accept and returns the same schema
+        of mapper outputs.
+    batch_rows : int or None
+        Rows in batches for mappers and reducers. Ignored if `mapper_batch_rows`
+        specified for mappers or `reducer_batch_rows` specified for reducers.
+        1024 by default.
+    mapper_dtypes : pd.Series or dict or None
+        Output dtypes of mapper stage.
+    mapper_index : pd.Index or None
+        Index of DataFrame returned by mappers.
+    mapper_batch_rows : int or None
+        Rows in batches for mappers. If specified, `batch_rows` will be ignored
+        for mappers.
+    reducer_dtypes : pd.Series or dict or None
+        Output dtypes of reducer stage.
+    reducer_index : pd.Index or None
+        Index of DataFrame returned by reducers.
+    reducer_batch_rows : int or None
+        Rows in batches for mappers. If specified, `batch_rows` will be ignored
+        for reducers.
+    ignore_index : bool
+        If true, indexes generated at mapper or reducer functions will be ignored.
+
+    Returns
+    -------
+    output: DataFrame
+        Result DataFrame after map and reduce.
+
+    Examples
+    --------
+
+    We first define a DataFrame with a column of several words.
+
+    >>> from collections import defaultdict
+    >>> import maxframe.dataframe as md
+    >>> from maxframe.udf import with_running_options
+    >>> df = pd.DataFrame(
+    >>>     {
+    >>>         "name": ["name key", "name", "key", "name", "key name"],
+    >>>         "id": [4, 2, 4, 3, 3],
+    >>>         "fid": [5.3, 3.5, 4.2, 2.2, 4.1],
+    >>>     }
+    >>> )
+
+    Then we write a mapper function which accepts batches in the DataFrame
+    and returns counts of words in every row.
+
+    >>> def mapper(batch):
+    >>>     word_to_count = defaultdict(lambda: 0)
+    >>>     for words in batch["name"]:
+    >>>         for w in words.split():
+    >>>             word_to_count[w] += 1
+    >>>     return pd.DataFrame(
+    >>>         [list(tp) for tp in word_to_count.items()], columns=["word", "count"]
+    >>>     )
+
+    After that we write a reducer function which aggregates records with
+    the same word. Running options such as CPU specifications can be supplied
+    as well.
+
+    >>> @with_running_options(cpu=2)
+    >>> class TestReducer:
+    >>>     def __init__(self):
+    >>>         self._word_to_count = defaultdict(lambda: 0)
+    >>>
+    >>>     def __call__(self, batch, end=False):
+    >>>         word = None
+    >>>         for _, row in batch.iterrows():
+    >>>             word = row.iloc[0]
+    >>>             self._word_to_count[row.iloc[0]] += row.iloc[1]
+    >>>         if end:
+    >>>             return pd.DataFrame(
+    >>>                 [[word, self._word_to_count[word]]], columns=["word", "count"]
+    >>>             )
+    >>>
+    >>>     def close(self):
+    >>>         # you can do several cleanups here
+    >>>         print("close")
+
+    Finally we can call `map_reduce` with mappers and reducers specified above.
+
+    >>> res = df.mf.map_reduce(
+    >>>     mapper,
+    >>>     TestReducer,
+    >>>     group_cols=["word"],
+    >>>     mapper_dtypes={"word": "str", "count": "int"},
+    >>>     mapper_index=pd.Index([0]),
+    >>>     reducer_dtypes={"word": "str", "count": "int"},
+    >>>     reducer_index=pd.Index([0]),
+    >>>     ignore_index=True,
+    >>> )
+    >>> res.execute().fetch()
+       word  count
+    0   key      3
+    1  name      4
+
+    See Also
+    --------
+    DataFrame.mf.apply_chunk, DataFrame.groupby.mf.apply_chunk
+    """
+    mapper_batch_rows = mapper_batch_rows or batch_rows
+    reducer_batch_rows = reducer_batch_rows or batch_rows
+
+    def check_arg(arg_type, locals_):
+        if locals_.get(arg_type) is not None:
+            return
+        for suffix in ("dtypes", "index"):
+            arg_name = f"{arg_type}_{suffix}"
+            if locals_.get(arg_name) is not None:
+                raise ValueError(f"Cannot specify {arg_name} when {arg_type} is None")
+
+    if mapper is None:
+        check_arg("mapper", locals())
+        mapped = df
+        group_cols = group_cols or df.dtypes.index
+        if combiner is not None:
+            raise ValueError("Combiner cannot be set when mapper is None")
+    else:
+        if combiner is not None:
+            mapper = _gen_combined_mapper(
+                mapper, combiner, group_cols, order_cols, ascending=ascending
+            )
+        mapped = df.mf.apply_chunk(
+            mapper,
+            batch_rows=mapper_batch_rows,
+            dtypes=mapper_dtypes,
+            output_type="dataframe",
+            index=mapper_index,
+        )
+        group_cols = group_cols or list(df.dtypes.index)
+
+    if reducer is None:
+        check_arg("reducer", locals())
+        res = mapped
+    else:
+        res = mapped.groupby(group_cols, group_keys=False)[
+            list(mapped.dtypes.index)
+        ].mf.apply_chunk(
+            reducer,
+            batch_rows=reducer_batch_rows,
+            dtypes=reducer_dtypes,
+            output_type="dataframe",
+            index=reducer_index,
+            order_cols=order_cols,
+            ascending=ascending,
+        )
+
+    if ignore_index:
+        return res.reset_index(drop=True)
+    return res

maxframe/dataframe/extensions/rebalance.py

@@ -0,0 +1,62 @@
+# Copyright 1999-2025 Alibaba Group Holding Ltd.
+#
+# 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.
+
+from ... import opcodes
+from ...serialization.serializables import Float64Field, Int64Field
+from ...tensor.extensions.rebalance import RebalanceMixin
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import validate_axis
+
+
+class DataFrameRebalance(RebalanceMixin, DataFrameOperatorMixin, DataFrameOperator):
+    _op_type_ = opcodes.REBALANCE
+
+    axis = Int64Field("axis")
+    factor = Float64Field("factor", default=None)
+    num_partitions = Int64Field("num_partitions")
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+
+def rebalance(df_or_series, axis=0, factor=None, num_partitions=None):
+    """
+    Make data more balanced across entire cluster.
+
+    Parameters
+    ----------
+    axis : int
+        The axis to rebalance.
+    factor : float
+        Specified so that number of chunks after balance is
+        total number of input chunks * factor.
+    num_partitions : int
+        Specified so the number of chunks are at most
+        num_partitions.
+
+    Returns
+    -------
+    Series or DataFrame
+        Result of DataFrame or Series after rebalanced.
+    """
+    axis = validate_axis(axis, df_or_series)
+    if num_partitions is None and factor is None:
+        raise ValueError("Need to specify num_partitions or factor")
+    if num_partitions is not None and factor is not None:
+        raise ValueError(
+            "num_partitions and factor cannot be specified at the same time"
+        )
+
+    op = DataFrameRebalance(axis=axis, factor=factor, num_partitions=num_partitions)
+    return op(df_or_series)

maxframe/dataframe/extensions/tests/test_apply_chunk.py

@@ -91,15 +91,22 @@ def test_apply_chunk_infer_dtypes_and_index(df1, df2, df3):
     assert result.index_value is df2.index_value
     assert result.dtypes.equals(df2.dtypes)
 
+    def process(data, param, k) -> pd.DataFrame[df2.dtypes]:
+        return data * param * k
+
+    result = df2.mf.apply_chunk(process, batch_rows=3, args=(4,), k=1)
+    assert result.index_value is df2.index_value
+    assert result.dtypes.equals(df2.dtypes)
+
     # mark functions
     from ....udf import with_python_requirements, with_resources
 
     @with_resources("empty.txt")
     @with_python_requirements("numpy")
-    def process(data, k):
+    def process(data, k) -> pd.DataFrame[df1.dtypes]:
         return data
 
-    result = df1.mf.apply_chunk(process, batch_rows=3, output_type="dataframe", k=1)
+    result = df1.mf.apply_chunk(process, batch_rows=3, k=1)
     assert result.index_value is df1.index_value
     assert result.dtypes.equals(df1.dtypes)
     assert isinstance(result.op.func, MarkedFunction)

maxframe/dataframe/extensions/tests/test_extensions.py

@@ -60,6 +60,40 @@ def df3():
     )
 
 
+@pytest.fixture
+def df4():
+    return DataFrame(
+        {
+            "name1": ["a", "b", "c", "d"],
+            "name2": ["a", "b", "c", "d"],
+            "num": [1, 2, 3, 4],
+            "kv": [
+                "k1=1.1,k2=3.1,k3=1.0",
+                "k1=7.1,k4=8.2",
+                "k5=1.2,k7=1.5",
+                "k3=1.1,k9=1",
+            ],
+            "vk": ["v1=1.1,v2=1.2", "v3=1.1,v4=1.2", "v5=1.1,v6=1.2", "v7=1.1,v8=1.2"],
+        }
+    )
+
+
+@pytest.fixture
+def df5():
+    return DataFrame(
+        {
+            "name1": ["name1", "name2", "name3", "name4", "name5"],
+            "name2": ["name1", "name2", "name3", "name4", "name5"],
+            "k1": [1.0, None, 7.1, None, None],
+            "k2": [3.0, 3.0, None, 1.2, 1.0],
+            "k3": [None, 5.1, None, 1.5, None],
+            "k5": [10.0, None, None, None, None],
+            "k7": [None, None, 8.2, None, None],
+            "k9": [None, None, None, None, 1.1],
+        }
+    )
+
+
 def test_flatmap(df1, df2, df3):
     def f(x, keys):
         if x["a"] in keys:
@@ -142,3 +176,23 @@ def test_flatjson():
         )
     with pytest.raises(ValueError):
         s1.mf.flatjson(["$.a"])
+
+
+def test_extract_kv(df4):
+    extract_kv_df = df4.mf.extract_kv(
+        columns=["kv", "vk"], kv_delim=",", item_delim="="
+    )
+    assert extract_kv_df.shape == (4, np.nan)
+    assert extract_kv_df.index_value.key == df4.index_value.key
+    with pytest.raises(ValueError):
+        df4.mf.extract_kv(columns=["name"])
+    with pytest.raises(ValueError):
+        df4.mf.extract_kv(columns=["num"])
+
+
+def test_collect_kv(df5):
+    collect_kv_df = df5.mf.collect_kv(columns=["k1", "k2", "k3", "k5", "k7", "k9"])
+    assert collect_kv_df.shape == (5, 3)
+    assert collect_kv_df.index_value.key == df5.index_value.key
+    with pytest.raises(ValueError):
+        df5.mf.collect_kv(columns=["num"])

maxframe/dataframe/extensions/tests/test_map_reduce.py

@@ -0,0 +1,135 @@
+# Copyright 1999-2025 Alibaba Group Holding Ltd.
+#
+# 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 functools
+from collections import defaultdict
+
+import pandas as pd
+import pytest
+
+from .... import dataframe as md
+from ...groupby.apply_chunk import GroupByApplyChunk
+from .. import DataFrameApplyChunk
+
+
+@pytest.fixture
+def df1():
+    return md.DataFrame(
+        {
+            "name": ["name key", "name", "key", "name", "key name"],
+            "id": [4, 2, 4, 3, 3],
+            "fid": [5.3, 3.5, 4.2, 2.2, 4.1],
+        }
+    )
+
+
+def test_map_reduce_with_map_only(df1):
+    func = functools.partial(lambda x: x)
+
+    with pytest.raises(ValueError):
+        df1.mf.map_reduce(func, reducer_dtypes={"col": "string"})
+
+    mapped = df1.mf.map_reduce(func)
+    assert isinstance(mapped.op, DataFrameApplyChunk)
+    assert mapped.op.func is func
+
+    map_combined = df1.mf.map_reduce(
+        func, combiner=func, mapper_dtypes=df1.dtypes, mapper_index=df1.index
+    )
+    assert isinstance(map_combined.op, DataFrameApplyChunk)
+    assert map_combined.op.func.__name__ == "CombinedMapper"
+
+
+def test_mapper_with_combiner(df1):
+    class BaseFunc:
+        def __init__(self):
+            self._word_to_count = defaultdict(lambda: 0)
+
+        def _collect_df(self):
+            word_to_count = self._word_to_count.copy()
+            self._word_to_count.clear()
+            return pd.DataFrame(
+                [list(tp) for tp in word_to_count.items()],
+                columns=["word", "count"],
+            )
+
+        def close(self):
+            print(f"Close {type(self)}")
+
+    class MapperCls(BaseFunc):
+        def __call__(self, batch, end=False):
+            for words in batch["name"]:
+                for w in words.split():
+                    self._word_to_count[w] += 1
+            if end:
+                return self._collect_df()
+
+    class CombinerCls(BaseFunc):
+        def __call__(self, batch, end=False):
+            for _, row in batch.iterrows():
+                self._word_to_count[row["word"]] = row["count"]
+            if end:
+                return self._collect_df()
+
+    map_combined = df1.mf.map_reduce(
+        MapperCls,
+        combiner=CombinerCls,
+        group_cols="word",
+        mapper_dtypes={"word": "str", "count": "int"},
+        mapper_index=df1.index,
+    )
+
+    raw = df1.op.data
+    combiner = map_combined.op.func()
+    ret1 = combiner(raw.iloc[:3], end=True)
+    ret2 = combiner(raw.iloc[3:], end=True)
+    close_ret = combiner.close()
+    expected1 = pd.DataFrame([["key", 2], ["name", 2]], columns=["word", "count"])
+    expected2 = pd.DataFrame([["key", 1], ["name", 2]], columns=["word", "count"])
+    assert close_ret is None
+    pd.testing.assert_frame_equal(ret1.reset_index(drop=True), expected1)
+    pd.testing.assert_frame_equal(ret2.reset_index(drop=True), expected2)
+
+
+def test_map_reduce_with_reduce_only(df1):
+    func = functools.partial(lambda x: x)
+
+    with pytest.raises(ValueError):
+        df1.mf.map_reduce(reducer=func, mapper_dtypes={"col": "string"})
+
+    reduced = df1.mf.map_reduce(reducer=func, group_cols="name")
+    assert isinstance(reduced.op, GroupByApplyChunk)
+    assert reduced.op.func is func
+    assert reduced.op.groupby_params["by"] == ["name"]
+
+
+def test_map_reduce_with_both_funcs(df1):
+    map_func = functools.partial(lambda x: x)
+
+    class ReducerCls:
+        def __call__(self, batch):
+            return batch
+
+    reduced = df1.mf.map_reduce(
+        mapper=map_func,
+        reducer=ReducerCls,
+        group_cols="name",
+        reducer_dtypes=df1.dtypes,
+        reducer_index=df1.index,
+    )
+    assert isinstance(reduced.op, GroupByApplyChunk)
+    assert reduced.op.func is ReducerCls
+    assert reduced.op.groupby_params["by"] == ["name"]
+    assert isinstance(reduced.inputs[0].op, DataFrameApplyChunk)
+    assert reduced.inputs[0].op.func is map_func

maxframe/dataframe/groupby/__init__.py

@@ -15,6 +15,8 @@
 # noinspection PyUnresolvedReferences
 from ..core import DataFrameGroupBy, GroupBy, SeriesGroupBy
 from .core import NamedAgg
+from .expanding import ExpandingGroupby
+from .rolling import RollingGroupby
 
 
 def _install():
@@ -24,12 +26,15 @@ def _install():
     from .apply import groupby_apply
     from .apply_chunk import df_groupby_apply_chunk
     from .core import groupby
-    from .cum import cumcount, cummax, cummin, cumprod, cumsum
+    from .expanding import cumcount, cummax, cummin, cumprod, cumsum, expanding
     from .extensions import DataFrameGroupByMaxFrameAccessor
     from .fill import bfill, ffill, fillna
     from .getitem import df_groupby_getitem
     from .head import head
+    from .rank import rank
+    from .rolling import rolling
     from .sample import groupby_sample
+    from .shift import shift
     from .transform import groupby_transform
 
     for cls in DATAFRAME_TYPE:
@@ -69,6 +74,12 @@ def _install():
         setattr(cls, "cumprod", cumprod)
         setattr(cls, "cumsum", cumsum)
 
+        setattr(cls, "expanding", expanding)
+        setattr(cls, "rolling", rolling)
+
+        setattr(cls, "shift", shift)
+        setattr(cls, "rank", rank)
+
         setattr(cls, "head", head)
 
         setattr(cls, "sample", groupby_sample)