maxframe 1.3.1__cp39-cp39-macosx_10_9_universal2.whl → 2.0.0b2__cp39-cp39-macosx_10_9_universal2.whl

maxframe/dataframe/groupby/apply.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import MutableMapping, Union
+
 import numpy as np
-import pandas as pd
 
 from ... import opcodes
 from ...core import OutputType
@@ -26,12 +27,13 @@ from ...serialization.serializables import (
     StringField,
     TupleField,
 )
-from ...utils import get_func_token, quiet_stdio, tokenize
+from ...udf import BuiltinFunction, MarkedFunction
+from ...utils import copy_if_possible, get_func_token, make_dtype, make_dtypes, tokenize
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import (
+    InferredDataFrameMeta,
     copy_func_scheduling_hints,
-    make_dtype,
-    make_dtypes,
+    infer_dataframe_return_value,
     parse_index,
     validate_output_types,
 )
@@ -56,6 +58,7 @@ class GroupByApply(
     args = TupleField("args", default_factory=tuple)
     kwds = DictField("kwds", default_factory=dict)
     maybe_agg = BoolField("maybe_agg", default=None)
+
     logic_key = StringField("logic_key", default=None)
     func_key = AnyField("func_key", default=None)
     need_clean_up_func = BoolField("need_clean_up_func", default=False)
@@ -65,6 +68,9 @@ class GroupByApply(
         if hasattr(self, "func"):
             copy_func_scheduling_hints(self.func, self)
 
+    def has_custom_code(self) -> bool:
+        return not isinstance(self.func, BuiltinFunction)
+
     def _update_key(self):
         values = [v for v in self._values_ if v is not self.func] + [
             get_func_token(self.func)
@@ -73,96 +79,66 @@ class GroupByApply(
         return self
 
     def _infer_df_func_returns(
-        self, in_groupby, in_df, dtypes=None, dtype=None, name=None, index=None
-    ):
-        index_value, output_type, new_dtypes = None, None, None
-
-        if self.output_types is not None and (dtypes is not None or dtype is not None):
-            ret_dtypes = dtypes if dtypes is not None else (dtype, name)
-            ret_index_value = parse_index(index) if index is not None else None
-            return ret_dtypes, ret_index_value
-
-        try:
-            infer_df = in_groupby.op.build_mock_groupby().apply(
-                self.func, *self.args, **self.kwds
-            )
-
-            if len(infer_df) <= 2:
-                # we create mock df with 4 rows, 2 groups
-                # if return df has 2 rows, we assume that
-                # it's an aggregation operation
-                self.maybe_agg = True
-
-            # todo return proper index when sort=True is implemented
-            index_value = parse_index(infer_df.index[:0], in_df.key, self.func)
-
-            # for backward compatibility
-            dtype = dtype if dtype is not None else dtypes
-            if isinstance(infer_df, pd.DataFrame):
-                output_type = output_type or OutputType.dataframe
-                new_dtypes = new_dtypes or infer_df.dtypes
-            elif isinstance(infer_df, pd.Series):
-                output_type = output_type or OutputType.series
-                new_dtypes = new_dtypes or (
-                    name or infer_df.name,
-                    dtype or infer_df.dtype,
-                )
-            else:
-                output_type = OutputType.series
-                new_dtypes = (name, dtype or pd.Series(infer_df).dtype)
-        except:  # noqa: E722  # nosec
-            pass
-
+        self, in_groupby, dtypes=None, dtype=None, name=None, index=None
+    ) -> InferredDataFrameMeta:
+        def infer_func(groupby_obj):
+            args = copy_if_possible(self.args)
+            kwds = copy_if_possible(self.kwds)
+            return groupby_obj.apply(self.func, *args, **kwds)
+
+        output_type = self.output_types[0] if self.output_types else None
+        inferred_meta = infer_dataframe_return_value(
+            in_groupby,
+            infer_func,
+            dtypes=dtypes,
+            dtype=dtype,
+            name=name,
+            index=index,
+            output_type=output_type,
+        )
         self.output_types = (
-            [output_type]
-            if not self.output_types and output_type
+            [inferred_meta.output_type]
+            if not self.output_types and inferred_meta.output_type
             else self.output_types
         )
-        dtypes = new_dtypes if dtypes is None else dtypes
-        index_value = index_value if index is None else parse_index(index)
-        return dtypes, index_value
+        self.maybe_agg = inferred_meta.maybe_agg
+        return inferred_meta
 
     def __call__(self, groupby, dtypes=None, dtype=None, name=None, index=None):
-        in_df = groupby
         if self.output_types and self.output_types[0] == OutputType.df_or_series:
             return self.new_df_or_series([groupby])
-        while in_df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
-            in_df = in_df.inputs[0]
-
-        with quiet_stdio():
-            dtypes, index_value = self._infer_df_func_returns(
-                groupby, in_df, dtypes, dtype=dtype, name=name, index=index
-            )
-        if index_value is None:
-            index_value = parse_index(None, (in_df.key, in_df.index_value.key))
-        for arg, desc in zip((self.output_types, dtypes), ("output_types", "dtypes")):
-            if arg is None:
-                raise TypeError(
-                    f"Cannot determine {desc} by calculating with enumerate data, "
-                    "please specify it as arguments"
-                )
 
+        inferred_meta = self._infer_df_func_returns(
+            groupby, dtypes=dtypes, dtype=dtype, name=name, index=index
+        )
+        inferred_meta.check_absence("output_type", "dtypes", "dtype")
         if self.output_types[0] == OutputType.dataframe:
-            new_shape = (np.nan, len(dtypes))
+            new_shape = (np.nan, len(inferred_meta.dtypes))
             return self.new_dataframe(
                 [groupby],
                 shape=new_shape,
-                dtypes=dtypes,
-                index_value=index_value,
-                columns_value=parse_index(dtypes.index, store_data=True),
+                dtypes=inferred_meta.dtypes,
+                index_value=inferred_meta.index_value,
+                columns_value=parse_index(inferred_meta.dtypes.index, store_data=True),
             )
         else:
-            name = name or dtypes[0]
-            dtype = dtype or dtypes[1]
             new_shape = (np.nan,)
             return self.new_series(
                 [groupby],
-                name=name,
+                name=inferred_meta.name,
                 shape=new_shape,
-                dtype=dtype,
-                index_value=index_value,
+                dtype=inferred_meta.dtype,
+                index_value=inferred_meta.index_value,
             )
 
+    @classmethod
+    def estimate_size(
+        cls, ctx: MutableMapping[str, Union[int, float]], op: "GroupByApply"
+    ) -> None:
+        if isinstance(op.func, MarkedFunction):
+            ctx[op.outputs[0].key] = float("inf")
+        super().estimate_size(ctx, op)
+
 
 def groupby_apply(
     groupby,

maxframe/dataframe/groupby/apply_chunk.py

@@ -0,0 +1,393 @@
+# 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 Any, Callable, Dict, List, MutableMapping, Tuple, Union
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...lib.version import parse as parse_version
+from ...serialization.serializables import (
+    DictField,
+    FunctionField,
+    Int32Field,
+    TupleField,
+)
+from ...udf import BuiltinFunction, MarkedFunction
+from ...utils import copy_if_possible
+from ..core import (
+    DATAFRAME_GROUPBY_TYPE,
+    GROUPBY_TYPE,
+    DataFrameGroupBy,
+    IndexValue,
+    SeriesGroupBy,
+)
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import (
+    InferredDataFrameMeta,
+    build_empty_df,
+    copy_func_scheduling_hints,
+    infer_dataframe_return_value,
+    make_column_list,
+    make_dtype,
+    make_dtypes,
+    parse_index,
+    validate_output_types,
+)
+
+_need_enforce_group_keys = parse_version(pd.__version__) < parse_version("1.5.0")
+
+
+class GroupByApplyChunk(DataFrameOperatorMixin, DataFrameOperator):
+    _op_type_ = opcodes.APPLY_CHUNK
+    _op_module_ = "dataframe.groupby"
+
+    func = FunctionField("func")
+    batch_rows = Int32Field("batch_rows", default=None)
+    args = TupleField("args", default=None)
+    kwargs = DictField("kwargs", default=None)
+
+    groupby_params = DictField("groupby_params", default=None)
+
+    def __init__(self, output_type=None, **kw):
+        if output_type:
+            kw["_output_types"] = [output_type]
+        super().__init__(**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, dtypes, dtype, name, index_value, element_wise):
+        # return dataframe
+        if self.output_types[0] == OutputType.dataframe:
+            dtypes = make_dtypes(dtypes)
+            # apply_chunk will use generate new range index for results
+            return self.new_dataframe(
+                [df],
+                shape=df.shape if element_wise else (np.nan, len(dtypes)),
+                index_value=index_value,
+                columns_value=parse_index(dtypes.index, store_data=True),
+                dtypes=dtypes,
+            )
+
+        # return series
+        return self.new_series(
+            [df], shape=(np.nan,), name=name, dtype=dtype, index_value=index_value
+        )
+
+    def _call_series(self, series, dtypes, dtype, name, index_value, element_wise):
+        if self.output_types[0] == OutputType.series:
+            shape = series.shape if element_wise else (np.nan,)
+            return self.new_series(
+                [series],
+                dtype=dtype,
+                shape=shape,
+                index_value=index_value,
+                name=name,
+            )
+
+        dtypes = make_dtypes(dtypes)
+        return self.new_dataframe(
+            [series],
+            shape=(np.nan, len(dtypes)),
+            index_value=index_value,
+            columns_value=parse_index(dtypes.index, store_data=True),
+            dtypes=dtypes,
+        )
+
+    def __call__(
+        self,
+        groupby: Union[DataFrameGroupBy, SeriesGroupBy],
+        dtypes: Union[Tuple[str, Any], Dict[str, Any]] = None,
+        dtype: Any = None,
+        name: Any = None,
+        output_type=None,
+        index=None,
+    ):
+        input_df = groupby.inputs[0]
+        if isinstance(input_df, GROUPBY_TYPE):
+            input_df = input_df.inputs[0]
+
+        # if skip_infer, directly build a frame
+        if self.output_types and self.output_types[0] == OutputType.df_or_series:
+            return self.new_df_or_series([input_df])
+
+        # infer return index and dtypes
+        inferred_meta = self._infer_batch_func_returns(
+            groupby,
+            output_type=output_type,
+            dtypes=dtypes,
+            dtype=dtype,
+            name=name,
+            index=index,
+        )
+
+        if inferred_meta.index_value is None:
+            inferred_meta.index_value = parse_index(
+                None, (groupby.key, groupby.index_value.key, self.func)
+            )
+        inferred_meta.check_absence("output_type", "dtypes", "dtype")
+
+        if isinstance(groupby, DATAFRAME_GROUPBY_TYPE):
+            return self._call_dataframe(
+                input_df,
+                dtypes=inferred_meta.dtypes,
+                dtype=inferred_meta.dtype,
+                name=inferred_meta.name,
+                index_value=inferred_meta.index_value,
+                element_wise=inferred_meta.elementwise,
+            )
+
+        return self._call_series(
+            input_df,
+            dtypes=inferred_meta.dtypes,
+            dtype=inferred_meta.dtype,
+            name=inferred_meta.name,
+            index_value=inferred_meta.index_value,
+            element_wise=inferred_meta.elementwise,
+        )
+
+    def _infer_batch_func_returns(
+        self,
+        input_groupby: Union[DataFrameGroupBy, SeriesGroupBy],
+        output_type: OutputType,
+        dtypes: Union[pd.Series, List[Any], Dict[str, Any]] = None,
+        dtype: Any = None,
+        name: Any = None,
+        index: Union[pd.Index, IndexValue] = None,
+        elementwise: bool = None,
+    ) -> InferredDataFrameMeta:
+        def infer_func(groupby_obj):
+            args = copy_if_possible(self.args or ())
+            kwargs = copy_if_possible(self.kwargs or {})
+
+            in_obj = input_groupby
+            while isinstance(in_obj, GROUPBY_TYPE):
+                in_obj = in_obj.inputs[0]
+
+            by_cols = make_column_list(groupby_params.get("by"), in_obj.dtypes) or []
+            if not groupby_params.get("selection"):
+                selection = [
+                    c for c in input_groupby.inputs[0].dtypes.index if c not in by_cols
+                ]
+                groupby_obj = groupby_obj[selection]
+            res = groupby_obj.apply(self.func, *args, **kwargs)
+            if _need_enforce_group_keys and groupby_params.get("group_keys"):
+                by_levels = (
+                    make_column_list(groupby_params.get("level"), in_obj.index.names)
+                    or []
+                )
+
+                input_df = input_groupby
+                while isinstance(input_df, GROUPBY_TYPE):
+                    input_df = input_df.inputs[0]
+
+                idx_df = res.index.to_frame()
+                if by_cols:
+                    idx_names = by_cols + list(res.index.names)
+                    mock_idx_df = build_empty_df(
+                        input_df.dtypes[by_cols], index=idx_df.index
+                    )
+                else:
+                    idx_names = by_levels + list(res.index.names)
+                    if len(in_obj.index.names) > 1:
+                        idx_dtypes = in_obj.index_value.value.dtypes
+                    else:
+                        idx_dtypes = pd.Series(
+                            [in_obj.index.dtype], index=[in_obj.index.name]
+                        )
+                    mock_idx_df = build_empty_df(
+                        idx_dtypes[by_levels], index=idx_df.index
+                    )
+                idx_df = pd.concat([mock_idx_df, idx_df], axis=1)
+                res.index = pd.MultiIndex.from_frame(idx_df, names=idx_names)
+            return res
+
+        groupby_params = input_groupby.op.groupby_params
+        inferred_meta = infer_dataframe_return_value(
+            input_groupby,
+            infer_func,
+            output_type=output_type,
+            dtypes=dtypes,
+            dtype=dtype,
+            name=name,
+            index=index,
+            elementwise=elementwise,
+        )
+
+        # merge specified and inferred index, dtypes, output_type
+        # elementwise used to decide shape
+        self.output_types = (
+            [inferred_meta.output_type]
+            if not self.output_types and inferred_meta.output_type
+            else self.output_types
+        )
+        if self.output_types:
+            inferred_meta.output_type = self.output_types[0]
+        inferred_meta.dtypes = dtypes if dtypes is not None else inferred_meta.dtypes
+        if index is not None:
+            inferred_meta.index_value = (
+                parse_index(index)
+                if index is not input_groupby.index_value
+                else input_groupby.index_value
+            )
+        else:
+            inferred_meta.index_value = inferred_meta.index_value
+        inferred_meta.elementwise = elementwise or inferred_meta.elementwise
+        return inferred_meta
+
+    @classmethod
+    def estimate_size(
+        cls,
+        ctx: MutableMapping[str, Union[int, float]],
+        op: "GroupByApplyChunk",
+    ) -> None:
+        if isinstance(op.func, MarkedFunction):
+            ctx[op.outputs[0].key] = float("inf")
+        super().estimate_size(ctx, op)
+
+
+def df_groupby_apply_chunk(
+    dataframe_groupby,
+    func: Union[str, Callable],
+    batch_rows=None,
+    dtypes=None,
+    dtype=None,
+    name=None,
+    output_type=None,
+    index=None,
+    skip_infer=False,
+    args=(),
+    **kwargs,
+):
+    """
+    Apply function `func` group-wise and combine the results together.
+    The pandas DataFrame given to the function is a chunk of the input
+    dataframe, consider as a batch rows.
+
+    The function passed to `apply` must take a dataframe as its first
+    argument and return a DataFrame, Series or scalar. `apply` will
+    then take care of combining the results back together into a single
+    dataframe or series. `apply` is therefore a highly flexible
+    grouping method.
+
+    Don't expect to receive all rows of the DataFrame in the function,
+    as it depends on the implementation of MaxFrame and the internal
+    running state of MaxCompute.
+
+    Parameters
+    ----------
+    func : callable
+        A callable that takes a dataframe as its first argument, and
+        returns a dataframe, a series or a scalar. In addition the
+        callable may take positional and keyword arguments.
+
+    batch_rows : int
+        Specify expected number of rows in a batch, as well as the len of
+        function input dataframe. When the remaining data is insufficient,
+        it may be less than this number.
+
+    output_type : {'dataframe', 'series'}, default None
+        Specify type of returned object. See `Notes` for more details.
+
+    dtypes : Series, default None
+        Specify dtypes of returned DataFrames. See `Notes` for more details.
+
+    dtype : numpy.dtype, default None
+        Specify dtype of returned Series. See `Notes` for more details.
+
+    name : str, default None
+        Specify name of returned Series. See `Notes` for more details.
+
+    index : Index, default None
+        Specify index of returned object. See `Notes` for more details.
+
+    skip_infer: bool, default False
+        Whether infer dtypes when dtypes or output_type is not specified.
+
+    args, kwargs : tuple and dict
+        Optional positional and keyword arguments to pass to `func`.
+
+    Returns
+    -------
+    applied : Series or DataFrame
+
+    See Also
+    --------
+    Series.apply : Apply a function to a Series.
+    DataFrame.apply : Apply a function to each row or column of a DataFrame.
+    DataFrame.mf.apply_chunk : Apply a function to row batches of a DataFrame.
+
+    Notes
+    -----
+    When deciding output dtypes and shape of the return value, MaxFrame will
+    try applying ``func`` onto a mock grouped object, and the apply call
+    may fail. When this happens, you need to specify the type of apply
+    call (DataFrame or Series) in output_type.
+
+    * For DataFrame output, you need to specify a list or a pandas Series
+      as ``dtypes`` of output DataFrame. ``index`` of output can also be
+      specified.
+    * For Series output, you need to specify ``dtype`` and ``name`` of
+      output Series.
+
+    MaxFrame adopts expected behavior of pandas>=3.0 by ignoring group columns
+    in user function input. If you still need a group column for your function
+    input, try selecting it right after `groupby` results, for instance,
+    ``df.groupby("A")[["A", "B", "C"]].mf.apply_batch(func)`` will pass data of
+    column A into ``func``.
+    """
+    if not isinstance(func, Callable):
+        raise TypeError("function must be a callable object")
+
+    if batch_rows is not None:
+        if not isinstance(batch_rows, int):
+            raise TypeError("batch_rows must be an integer")
+        elif batch_rows <= 0:
+            raise ValueError("batch_rows must be greater than 0")
+
+    if dtype is not None:
+        dtype = make_dtype(dtype)
+
+    output_types = kwargs.pop("output_types", None)
+    object_type = kwargs.pop("object_type", None)
+    output_types = validate_output_types(
+        output_type=output_type, output_types=output_types, object_type=object_type
+    )
+    output_type = output_types[0] if output_types else None
+    if skip_infer and output_type is None:
+        output_type = OutputType.df_or_series
+
+    # bind args and kwargs
+    op = GroupByApplyChunk(
+        func=func,
+        batch_rows=batch_rows,
+        output_type=output_type,
+        args=args,
+        kwargs=kwargs,
+        groupby_params=dataframe_groupby.op.groupby_params,
+    )
+
+    return op(
+        dataframe_groupby,
+        dtypes=dtypes,
+        dtype=dtype,
+        name=name,
+        index=index,
+        output_type=output_type,
+    )