maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/dataframe/extensions/flatmap.py

@@ -0,0 +1,329 @@
+# 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 typing import Callable, MutableMapping, Union
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import (
+    BoolField,
+    DictField,
+    FunctionField,
+    TupleField,
+)
+from ...udf import BuiltinFunction, MarkedFunction
+from ...utils import make_dtypes
+from ..core import DataFrame
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import copy_func_scheduling_hints, gen_unknown_index_value, parse_index
+
+
+class DataFrameFlatMapOperator(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.FLATMAP
+
+    func = FunctionField("func")
+    raw = BoolField("raw", default=False)
+    args = TupleField("args", default=())
+    kwargs = DictField("kwargs", default={})
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+        if hasattr(self, "func"):
+            copy_func_scheduling_hints(self.func, self)
+
+    def has_custom_code(self) -> bool:
+        return not isinstance(self.func, BuiltinFunction)
+
+    def _call_dataframe(self, df: DataFrame, dtypes: pd.Series):
+        dtypes = make_dtypes(dtypes)
+        index_value = gen_unknown_index_value(
+            df.index_value,
+            (df.key, df.index_value.key, self.func),
+            normalize_range_index=True,
+        )
+        return self.new_dataframe(
+            [df],
+            shape=(np.nan, len(dtypes)),
+            index_value=index_value,
+            columns_value=parse_index(dtypes.index, store_data=True),
+            dtypes=dtypes,
+        )
+
+    def _call_series_or_index(self, series, dtypes=None):
+        index_value = gen_unknown_index_value(
+            series.index_value,
+            (series.key, series.index_value.key, self.func),
+            normalize_range_index=True,
+        )
+
+        if self.output_types[0] == OutputType.series:
+            name, dtype = dtypes
+            return self.new_series(
+                [series],
+                dtype=dtype,
+                shape=(np.nan,),
+                index_value=index_value,
+                name=name,
+            )
+
+        dtypes = make_dtypes(dtypes)
+        columns_value = parse_index(dtypes.index, store_data=True)
+        return self.new_dataframe(
+            [series],
+            shape=(np.nan, len(dtypes)),
+            index_value=index_value,
+            columns_value=columns_value,
+            dtypes=dtypes,
+        )
+
+    def __call__(
+        self,
+        df_or_series,
+        dtypes=None,
+        output_type=None,
+    ):
+        if df_or_series.op.output_types[0] == OutputType.dataframe:
+            return self._call_dataframe(df_or_series, dtypes=dtypes)
+        else:
+            return self._call_series_or_index(df_or_series, dtypes=dtypes)
+
+    @classmethod
+    def estimate_size(
+        cls, ctx: MutableMapping[str, Union[int, float]], op: "DataFrameFlatMapOperator"
+    ) -> None:
+        if isinstance(op.func, MarkedFunction):
+            ctx[op.outputs[0].key] = float("inf")
+        super().estimate_size(ctx, op)
+
+
+def df_flatmap(dataframe, func: Callable, dtypes=None, raw=False, args=(), **kwargs):
+    """
+    Apply the given function to each row and then flatten results. Use this method if your transformation returns
+    multiple rows for each input row.
+
+    This function applies a transformation to each row of the DataFrame, where the transformation can return zero
+    or multiple values, effectively flattening Python generators, list-like collections, and DataFrames.
+
+    Parameters
+    ----------
+    func : Callable
+        Function to apply to each row of the DataFrame. It should accept a Series (or an array if `raw=True`)
+        representing a row and return a list or iterable of values.
+
+    dtypes : Series, dict or list
+        Specify dtypes of returned DataFrame.
+
+    raw : bool, default False
+        Determines if the row is passed as a Series or as a numpy array:
+
+        * ``False`` : passes each row as a Series to the function.
+        * ``True`` : the passed function will receive numpy array objects instead.
+
+    args : tuple
+        Positional arguments to pass to `func`.
+
+    **kwargs
+        Additional keyword arguments to pass as keywords arguments to `func`.
+
+    Returns
+    -------
+    DataFrame
+        Return DataFrame with specified `dtypes`.
+
+    Notes
+    -----
+    The ``func`` must return an iterable of values for each input row. The index of the resulting DataFrame will be
+    repeated based on the number of output rows generated by `func`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
+    >>> df.execute()
+       A  B
+    0  1  4
+    1  2  5
+    2  3  6
+
+    Define a function that takes a number and returns a list of two numbers:
+
+    >>> def generate_values_array(row):
+    ...     return [row['A'] * 2, row['B'] * 3]
+
+    Define a function that takes a row and return two rows and two columns:
+
+    >>> def generate_values_in_generator(row):
+    ...     yield [row[0] * 2, row[1] * 4]
+    ...     yield [row[0] * 3, row[1] * 5]
+
+    Which equals to the following function return a dataframe:
+
+    >>> def generate_values_in_dataframe(row):
+    ...     return pd.DataFrame([[row[0] * 2, row[1] * 4], [row[0] * 3, row[1] * 5]])
+
+    Specify `dtypes` with a function which returns a DataFrame:
+
+    >>> df.mf.flatmap(generate_values_array, dtypes=pd.Series({'A': 'int'})).execute()
+            A
+        0   2
+        0  12
+        1   4
+        1  15
+        2   6
+        2  18
+
+    Specify raw=True to pass input row as array:
+
+    >>> df.mf.flatmap(generate_values_in_generator, dtypes={"A": "int", "B": "int"}, raw=True).execute()
+           A   B
+        0  2  16
+        0  3  20
+        1  4  20
+        1  6  25
+        2  6  24
+        2  9  30
+    """
+    if dtypes is None or len(dtypes) == 0:
+        raise TypeError(
+            "Cannot determine {dtypes} by calculating with enumerate data, "
+            "please specify it as arguments"
+        )
+
+    if not isinstance(func, Callable):
+        raise TypeError("function must be a callable object")
+
+    output_types = [OutputType.dataframe]
+    op = DataFrameFlatMapOperator(
+        func=func, raw=raw, output_types=output_types, args=args, kwargs=kwargs
+    )
+    return op(
+        dataframe,
+        dtypes=dtypes,
+    )
+
+
+def series_flatmap(
+    series, func: Callable, dtypes=None, dtype=None, name=None, args=(), **kwargs
+):
+    """
+    Apply the given function to each row and then flatten results. Use this method if your transformation returns
+    multiple rows for each input row.
+
+    This function applies a transformation to each element of the Series, where the transformation can return zero
+     or multiple values, effectively flattening Python generator, list-liked collections and DataFrame.
+
+    Parameters
+    ----------
+    func : Callable
+        Function to apply to each element of the Series. It should accept a scalar value
+        (or an array if ``raw=True``) and return a list or iterable of values.
+
+    dtypes : Series, default None
+        Specify dtypes of returned DataFrame. Can't work with dtype.
+
+    dtype : numpy.dtype, default None
+        Specify dtype of returned Series. Can't work with dtypes.
+
+    name : str, default None
+        Specify name of the returned Series.
+
+    args : tuple
+        Positional arguments to pass to ``func``.
+
+    **kwargs
+        Additional keyword arguments to pass as keywords arguments to ``func``.
+
+    Returns
+    -------
+    DataFrame or Series
+        Result of DataFrame when dtypes specified, else Series.
+
+    Notes
+    -----
+    The ``func`` must return an iterable of values for each input element. If ``dtypes`` is specified,
+    `flatmap` will return a DataFrame, if ``dtype`` and ``name`` is specified, a Series will be returned.
+
+    The index of the resulting DataFrame/Series will be repeated based on the number of output rows generated
+    by ``func``.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
+    >>> df.execute()
+       A  B
+    0  1  4
+    1  2  5
+    2  3  6
+
+    Define a function that takes a number and returns a list of two numbers:
+
+    >>> def generate_values_array(x):
+    ...     return [x * 2, x * 3]
+
+    Specify ``dtype`` with a function which returns list to return more elements as a Series:
+
+    >>> df['A'].mf.flatmap(generate_values_array, dtype="int", name="C").execute()
+        0    2
+        0    3
+        1    4
+        1    6
+        2    6
+        2    9
+        Name: C, dtype: int64
+
+    Specify ``dtypes`` to return multi columns as a DataFrame:
+
+
+    >>> def generate_values_in_generator(x):
+    ...     yield pd.Series([x * 2, x * 4])
+    ...     yield pd.Series([x * 3, x * 5])
+
+    >>> df['A'].mf.flatmap(generate_values_in_generator, dtypes={"A": "int", "B": "int"}).execute()
+           A   B
+        0  2   4
+        0  3   5
+        1  4   8
+        1  6  10
+        2  6  12
+        2  9  15
+    """
+
+    if dtypes is not None and dtype is not None:
+        raise ValueError("Both dtypes and dtype cannot be specified at the same time.")
+
+    dtypes = (name, dtype) if dtype is not None else dtypes
+    if dtypes is None:
+        raise TypeError(
+            "Cannot determine {dtypes} or {dtype} by calculating with enumerate data, "
+            "please specify it as arguments"
+        )
+
+    if not isinstance(func, Callable):
+        raise TypeError("function must be a callable object")
+
+    output_type = OutputType.series if dtype is not None else OutputType.dataframe
+
+    op = DataFrameFlatMapOperator(
+        func=func, raw=False, output_types=[output_type], args=args, kwargs=kwargs
+    )
+    return op(
+        series,
+        dtypes=dtypes,
+    )

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)