maxframe 0.1.0b5__cp38-cp38-macosx_11_0_arm64.whl

maxframe/dataframe/groupby/core.py

@@ -0,0 +1,179 @@
+# Copyright 1999-2024 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 collections import namedtuple
+
+import pandas as pd
+
+from ... import opcodes
+from ...core import ENTITY_TYPE, Entity, OutputType
+from ...core.operator import MapReduceOperator
+from ...serialization.serializables import AnyField, BoolField, Int32Field
+from ...utils import lazy_import, no_default, pd_release_version
+from ..core import SERIES_TYPE
+from ..initializer import Series as asseries
+from ..operators import DataFrameOperatorMixin
+from ..utils import build_df, build_series, parse_index
+
+cudf = lazy_import("cudf")
+
+_GROUP_KEYS_NO_DEFAULT = pd_release_version >= (1, 5, 0)
+_default_group_keys = no_default if _GROUP_KEYS_NO_DEFAULT else True
+
+
+NamedAgg = namedtuple("NamedAgg", ["column", "aggfunc"])
+
+
+class DataFrameGroupByOperator(MapReduceOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.GROUPBY
+
+    by = AnyField(
+        "by",
+        default=None,
+        on_serialize=lambda x: x.data if isinstance(x, Entity) else x,
+    )
+    level = AnyField("level", default=None)
+    as_index = BoolField("as_index", default=None)
+    sort = BoolField("sort", default=None)
+    group_keys = BoolField("group_keys", default=None)
+
+    shuffle_size = Int32Field("shuffle_size", default=None)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+        if output_types:
+            if output_types[0] in (
+                OutputType.dataframe,
+                OutputType.dataframe_groupby,
+            ):
+                output_types = [OutputType.dataframe_groupby]
+            elif output_types[0] == OutputType.series:
+                output_types = [OutputType.series_groupby]
+            self.output_types = output_types
+
+    @property
+    def is_dataframe_obj(self):
+        return self.output_types[0] in (
+            OutputType.dataframe_groupby,
+            OutputType.dataframe,
+        )
+
+    @property
+    def groupby_params(self):
+        return dict(
+            by=self.by,
+            level=self.level,
+            as_index=self.as_index,
+            sort=self.sort,
+            group_keys=self.group_keys,
+        )
+
+    def build_mock_groupby(self, **kwargs):
+        in_df = self.inputs[0]
+        if self.is_dataframe_obj:
+            mock_obj = build_df(
+                in_df, size=[2, 2], fill_value=[1, 2], ensure_string=True
+            )
+        else:
+            mock_obj = build_series(
+                in_df,
+                size=[2, 2],
+                fill_value=[1, 2],
+                name=in_df.name,
+                ensure_string=True,
+            )
+
+        new_kw = self.groupby_params
+        new_kw.update(kwargs)
+        if new_kw.get("level"):
+            new_kw["level"] = 0
+        if isinstance(new_kw["by"], list):
+            new_by = []
+            for v in new_kw["by"]:
+                if isinstance(v, ENTITY_TYPE):
+                    build_fun = build_df if v.ndim == 2 else build_series
+                    mock_by = pd.concat(
+                        [
+                            build_fun(v, size=2, fill_value=1, name=v.name),
+                            build_fun(v, size=2, fill_value=2, name=v.name),
+                        ]
+                    )
+                    new_by.append(mock_by)
+                else:
+                    new_by.append(v)
+            new_kw["by"] = new_by
+        return mock_obj.groupby(**new_kw)
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        inputs_iter = iter(self._inputs[1:])
+        if len(inputs) > 1:
+            by = []
+            for k in self.by:
+                if isinstance(k, SERIES_TYPE):
+                    by.append(next(inputs_iter))
+                else:
+                    by.append(k)
+            self.by = by
+
+    def __call__(self, df):
+        params = df.params.copy()
+        params["index_value"] = parse_index(None, df.key, df.index_value.key)
+        if df.ndim == 2:
+            if isinstance(self.by, list):
+                index, types = [], []
+                for k in self.by:
+                    if isinstance(k, SERIES_TYPE):
+                        index.append(k.name)
+                        types.append(k.dtype)
+                    elif k in df.dtypes:
+                        index.append(k)
+                        types.append(df.dtypes[k])
+                    else:
+                        raise KeyError(k)
+                params["key_dtypes"] = pd.Series(types, index=index)
+
+        inputs = [df]
+        if isinstance(self.by, list):
+            for k in self.by:
+                if isinstance(k, SERIES_TYPE):
+                    inputs.append(k)
+
+        return self.new_tileable(inputs, **params)
+
+
+def groupby(
+    df, by=None, level=None, as_index=True, sort=True, group_keys=_default_group_keys
+):
+    if not as_index and df.op.output_types[0] == OutputType.series:
+        raise TypeError("as_index=False only valid with DataFrame")
+
+    output_types = (
+        [OutputType.dataframe_groupby] if df.ndim == 2 else [OutputType.series_groupby]
+    )
+    if isinstance(by, (SERIES_TYPE, pd.Series)):
+        if isinstance(by, pd.Series):
+            by = asseries(by)
+        by = [by]
+    elif df.ndim > 1 and by is not None and not isinstance(by, list):
+        by = [by]
+    op = DataFrameGroupByOperator(
+        by=by,
+        level=level,
+        as_index=as_index,
+        sort=sort,
+        group_keys=group_keys if group_keys is not no_default else None,
+        output_types=output_types,
+    )
+    return op(df)

maxframe/dataframe/groupby/cum.py

@@ -0,0 +1,124 @@
+# Copyright 1999-2024 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 pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import AnyField, BoolField
+from ...utils import lazy_import
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import parse_index, validate_axis
+
+cudf = lazy_import("cudf")
+
+
+class GroupByCumReductionOperator(DataFrameOperatorMixin, DataFrameOperator):
+    _op_module_ = "dataframe.groupby"
+
+    axis = AnyField("axis", default=None)
+    ascending = BoolField("ascending", default=None)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+    def _calc_out_dtypes(self, in_groupby):
+        mock_groupby = in_groupby.op.build_mock_groupby()
+        func_name = getattr(self, "_func_name")
+
+        if func_name == "cumcount":
+            result_df = mock_groupby.cumcount(ascending=self.ascending)
+        else:
+            result_df = getattr(mock_groupby, func_name)(axis=self.axis)
+
+        if isinstance(result_df, pd.DataFrame):
+            self.output_types = [OutputType.dataframe]
+            return result_df.dtypes
+        else:
+            self.output_types = [OutputType.series]
+            return result_df.name, result_df.dtype
+
+    def __call__(self, groupby):
+        in_df = groupby
+        while in_df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
+            in_df = in_df.inputs[0]
+
+        self.axis = validate_axis(self.axis or 0, in_df)
+
+        out_dtypes = self._calc_out_dtypes(groupby)
+
+        kw = in_df.params.copy()
+        kw["index_value"] = parse_index(pd.RangeIndex(-1), groupby.key)
+        if self.output_types[0] == OutputType.dataframe:
+            kw.update(
+                dict(
+                    columns_value=parse_index(out_dtypes.index, store_data=True),
+                    dtypes=out_dtypes,
+                    shape=(groupby.shape[0], len(out_dtypes)),
+                )
+            )
+        else:
+            name, dtype = out_dtypes
+            kw.update(dtype=dtype, name=name, shape=(groupby.shape[0],))
+        return self.new_tileable([groupby], **kw)
+
+
+class GroupByCummin(GroupByCumReductionOperator):
+    _op_type_ = opcodes.CUMMIN
+    _func_name = "cummin"
+
+
+class GroupByCummax(GroupByCumReductionOperator):
+    _op_type_ = opcodes.CUMMAX
+    _func_name = "cummax"
+
+
+class GroupByCumsum(GroupByCumReductionOperator):
+    _op_type_ = opcodes.CUMSUM
+    _func_name = "cumsum"
+
+
+class GroupByCumprod(GroupByCumReductionOperator):
+    _op_type_ = opcodes.CUMPROD
+    _func_name = "cumprod"
+
+
+class GroupByCumcount(GroupByCumReductionOperator):
+    _op_type_ = opcodes.CUMCOUNT
+    _func_name = "cumcount"
+
+
+def cumcount(groupby, ascending: bool = True):
+    op = GroupByCumcount(ascending=ascending)
+    return op(groupby)
+
+
+def cummin(groupby, axis=0):
+    op = GroupByCummin(axis=axis)
+    return op(groupby)
+
+
+def cummax(groupby, axis=0):
+    op = GroupByCummax(axis=axis)
+    return op(groupby)
+
+
+def cumprod(groupby, axis=0):
+    op = GroupByCumprod(axis=axis)
+    return op(groupby)
+
+
+def cumsum(groupby, axis=0):
+    op = GroupByCumsum(axis=axis)
+    return op(groupby)

maxframe/dataframe/groupby/fill.py

@@ -0,0 +1,141 @@
+# Copyright 1999-2024 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 pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import AnyField, DictField, Int64Field, StringField
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import parse_index
+
+
+class GroupByFillOperator(DataFrameOperator, DataFrameOperatorMixin):
+    _op_module_ = "dataframe.groupby"
+
+    value = AnyField("value", default=None)
+    method = StringField("method", default=None)
+    axis = AnyField("axis", default=0)
+    limit = Int64Field("limit", default=None)
+    downcast = DictField("downcast", default=None)
+
+    def _calc_out_dtypes(self, in_groupby):
+        mock_groupby = in_groupby.op.build_mock_groupby()
+        func_name = getattr(self, "_func_name")
+
+        if func_name == "fillna":
+            result_df = mock_groupby.fillna(
+                value=self.value,
+                method=self.method,
+                axis=self.axis,
+                limit=self.limit,
+                downcast=self.downcast,
+            )
+        else:
+            result_df = getattr(mock_groupby, func_name)(limit=self.limit)
+
+        if isinstance(result_df, pd.DataFrame):
+            self.output_types = [OutputType.dataframe]
+            return result_df.dtypes
+        else:
+            self.output_types = [OutputType.series]
+            return result_df.name, result_df.dtype
+
+    def __call__(self, groupby):
+        in_df = groupby
+        while in_df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
+            in_df = in_df.inputs[0]
+        out_dtypes = self._calc_out_dtypes(groupby)
+
+        kw = in_df.params.copy()
+        kw["index_value"] = parse_index(pd.RangeIndex(-1), groupby.key)
+        if self.output_types[0] == OutputType.dataframe:
+            kw.update(
+                dict(
+                    columns_value=parse_index(out_dtypes.index, store_data=True),
+                    dtypes=out_dtypes,
+                    shape=(groupby.shape[0], len(out_dtypes)),
+                )
+            )
+        else:
+            name, dtype = out_dtypes
+            kw.update(dtype=dtype, name=name, shape=(groupby.shape[0],))
+        return self.new_tileable([groupby], **kw)
+
+
+class GroupByFFill(GroupByFillOperator):
+    _op_type_ = opcodes.FILL_NA
+    _func_name = "ffill"
+
+
+class GroupByBFill(GroupByFillOperator):
+    _op_type = opcodes.FILL_NA
+    _func_name = "bfill"
+
+
+class GroupByFillNa(GroupByFillOperator):
+    _op_type = opcodes.FILL_NA
+    _func_name = "fillna"
+
+
+def ffill(groupby, limit=None):
+    """
+    Forward fill the values.
+
+    limit:  int, default None
+            Limit number of values to fill
+
+    return: Series or DataFrame
+    """
+    op = GroupByFFill(limit=limit)
+    return op(groupby)
+
+
+def bfill(groupby, limit=None):
+    """
+    Backward fill the values.
+
+    limit:  int, default None
+            Limit number of values to fill
+
+    return: Series or DataFrame
+    """
+    op = GroupByBFill(limit=limit)
+    return op(groupby)
+
+
+def fillna(groupby, value=None, method=None, axis=None, limit=None, downcast=None):
+    """
+    Fill NA/NaN values using the specified method
+
+    value:  scalar, dict, Series, or DataFrame
+            Value to use to fill holes (e.g. 0), alternately a dict/Series/DataFrame
+            of values specifying which value to use for each index (for a Series) or
+            column (for a DataFrame). Values not in the dict/Series/DataFrame
+            will not be filled. This value cannot be a list.
+    method: {'backfill','bfill','ffill',None}, default None
+    axis:   {0 or 'index', 1 or 'column'}
+    limit:  int, default None
+            If method is specified, this is the maximum number of consecutive
+            NaN values to forward/backward fill
+    downcast:   dict, default None
+                A dict of item->dtype of what to downcast if possible,
+                or the string ‘infer’ which will try to downcast to an appropriate equal type
+
+    return: DataFrame or None
+    """
+    op = GroupByFillNa(
+        value=value, method=method, axis=axis, limit=limit, downcast=downcast
+    )
+    return op(groupby)

maxframe/dataframe/groupby/getitem.py

@@ -0,0 +1,92 @@
+# Copyright 1999-2024 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 collections.abc import Iterable
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import AnyField
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import parse_index
+
+
+class GroupByIndex(DataFrameOperatorMixin, DataFrameOperator):
+    _op_type_ = opcodes.INDEX
+    _op_module_ = "dataframe.groupby"
+
+    selection = AnyField("selection", default=None)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+    @property
+    def groupby_params(self):
+        params = self.inputs[0].op.groupby_params
+        params["selection"] = self.selection
+        return params
+
+    def build_mock_groupby(self, **kwargs):
+        groupby_op = self.inputs[0].op
+        return groupby_op.build_mock_groupby(**kwargs)[self.selection]
+
+    def __call__(self, groupby):
+        indexed = groupby.op.build_mock_groupby()[self.selection]
+
+        if indexed.ndim == 1:
+            self.output_types = [OutputType.series_groupby]
+            params = dict(
+                shape=(groupby.shape[0],),
+                name=self.selection,
+                dtype=groupby.dtypes[self.selection],
+                index_value=groupby.index_value,
+                key_dtypes=groupby.key_dtypes,
+            )
+        else:
+            self.output_types = [OutputType.dataframe_groupby]
+
+            if isinstance(self.selection, Iterable) and not isinstance(
+                self.selection, str
+            ):
+                item_list = list(self.selection)
+            else:
+                item_list = [self.selection]
+
+            params = groupby.params.copy()
+            params["dtypes"] = new_dtypes = groupby.dtypes[item_list]
+            params["selection"] = self.selection
+            params["shape"] = (groupby.shape[0], len(item_list))
+            params["columns_value"] = parse_index(new_dtypes.index, store_data=True)
+
+        return self.new_tileable([groupby], **params)
+
+
+def df_groupby_getitem(df_groupby, item):
+    try:
+        hash(item)
+        hashable = True
+    except TypeError:
+        hashable = False
+
+    if hashable and item in df_groupby.dtypes:
+        output_types = [OutputType.series_groupby]
+    elif isinstance(item, Iterable) and all(it in df_groupby.dtypes for it in item):
+        output_types = [OutputType.dataframe_groupby]
+    else:
+        raise NameError(f"Cannot slice groupby with {item!r}")
+
+    if df_groupby.selection:
+        raise IndexError(f"Column(s) {df_groupby.selection!r} already selected")
+
+    op = GroupByIndex(selection=item, output_types=output_types)
+    return op(df_groupby)

maxframe/dataframe/groupby/head.py

@@ -0,0 +1,105 @@
+# Copyright 1999-2024 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 numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import BoolField, DictField, Int64Field
+from ...utils import pd_release_version
+from ..core import IndexValue
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import parse_index
+
+_pandas_enable_negative = pd_release_version >= (1, 4, 0)
+
+
+class GroupByHead(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.GROUPBY_HEAD
+    _op_module_ = "dataframe.groupby"
+
+    row_count = Int64Field("row_count", default=5)
+    groupby_params = DictField("groupby_params", default=dict())
+    enable_negative = BoolField("enable_negative", default=_pandas_enable_negative)
+
+    def __call__(self, groupby):
+        df = groupby
+        while df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
+            df = df.inputs[0]
+
+        selection = groupby.op.groupby_params.pop("selection", None)
+        if df.ndim > 1 and selection:
+            if isinstance(selection, tuple) and selection not in df.dtypes:
+                selection = list(selection)
+
+            result_df = df[selection]
+        else:
+            result_df = df
+
+        self._output_types = (
+            [OutputType.dataframe] if result_df.ndim == 2 else [OutputType.series]
+        )
+
+        params = result_df.params
+        params["shape"] = (np.nan,) + result_df.shape[1:]
+        if isinstance(df.index_value.value, IndexValue.RangeIndex):
+            params["index_value"] = parse_index(pd.RangeIndex(-1), df.key)
+
+        return self.new_tileable([df], **params)
+
+
+def head(groupby, n=5):
+    """
+    Return first n rows of each group.
+
+    Similar to ``.apply(lambda x: x.head(n))``, but it returns a subset of rows
+    from the original Series or DataFrame with original index and order preserved
+    (``as_index`` flag is ignored).
+
+    Does not work for negative values of `n`.
+
+    Returns
+    -------
+    Series or DataFrame
+
+    See Also
+    --------
+    Series.groupby
+    DataFrame.groupby
+
+    Examples
+    --------
+
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame([[1, 2], [1, 4], [5, 6]],
+    ...                   columns=['A', 'B'])
+    >>> df.groupby('A').head(1).execute()
+       A  B
+    0  1  2
+    2  5  6
+    >>> df.groupby('A').head(-1).execute()
+    Empty DataFrame
+    Columns: [A, B]
+    Index: []
+    """
+    groupby_params = groupby.op.groupby_params.copy()
+    groupby_params.pop("as_index", None)
+
+    op = GroupByHead(
+        row_count=n,
+        groupby_params=groupby_params,
+        enable_negative=_pandas_enable_negative,
+    )
+    return op(groupby)