maxframe 0.1.0b5__cp311-cp311-macosx_11_0_arm64.whl

maxframe/dataframe/datastore/to_odps.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# 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 logging
+from typing import List, Optional, Union
+
+from odps.models import Table as ODPSTable
+
+from ... import opcodes
+from ...config import options
+from ...core import OutputType
+from ...serialization.serializables import (
+    BoolField,
+    FieldTypes,
+    Int64Field,
+    ListField,
+    SeriesField,
+    StringField,
+)
+from ...typing_ import TileableType
+from ..core import DataFrame  # noqa: F401
+from ..utils import parse_index
+from .core import DataFrameDataStore
+
+logger = logging.getLogger(__name__)
+
+
+class DataFrameToODPSTable(DataFrameDataStore):
+    _op_type_ = opcodes.TO_ODPS_TABLE
+
+    dtypes = SeriesField("dtypes")
+
+    table_name = StringField("table_name", default=None)
+    partition_spec = StringField("partition_spec", default=None)
+    partition_columns = ListField("partition_columns", FieldTypes.string, default=None)
+    overwrite = BoolField("overwrite", default=None)
+    write_batch_size = Int64Field("write_batch_size", default=None)
+    unknown_as_string = BoolField("unknown_as_string", default=None)
+    index = BoolField("index", default=True)
+    index_label = ListField("index_label", FieldTypes.string, default=None)
+    lifecycle = Int64Field("lifecycle", default=None)
+
+    def __init__(self, **kw):
+        super().__init__(_output_types=[OutputType.dataframe], **kw)
+
+    def __call__(self, x):
+        shape = (0,) * len(x.shape)
+        index_value = parse_index(x.index_value.to_pandas()[:0], x.key, "index")
+        columns_value = parse_index(
+            x.columns_value.to_pandas()[:0], x.key, "columns", store_data=True
+        )
+        return self.new_dataframe(
+            [x],
+            shape=shape,
+            dtypes=x.dtypes[:0],
+            index_value=index_value,
+            columns_value=columns_value,
+        )
+
+
+def to_odps_table(
+    df: TileableType,
+    table: Union[ODPSTable, str],
+    partition: Optional[str] = None,
+    partition_col: Union[None, str, List[str]] = None,
+    overwrite: bool = False,
+    unknown_as_string: Optional[bool] = None,
+    index: bool = True,
+    index_label: Union[None, str, List[str]] = None,
+    lifecycle: Optional[int] = None,
+):
+    """
+    Write DataFrame object into a MaxCompute (ODPS) table.
+
+    You need to provide the name of the table to write to. If you want to store
+    data into a specific partitioned of a table, argument `partition` can be used.
+    You can also use `partition_col` to specify DataFrame columns as partition
+    columns, and data in the DataFrame will be grouped by these columns and
+    inserted into partitions the values of these columns.
+
+    If the table does not exist, `to_odps_table` will create one.
+
+    Column names for indexes is determined by `index_label` argument. If the
+    argument is absent, names of the levels is used if they are not None, or
+    default names will be used. The default name for indexes with only one level
+    will be `index`, and for indexes with multiple levels, the name will be
+    `level_x` while x is the index of the level.
+
+    Parameters
+    ----------
+    table: str
+        Name ot the table to write DataFrame into
+    partition: Optional[str]
+        Spec of the partition to write to, can be 'pt1=xxx,pt2=yyy'
+    partition_col: Union[None, str, List[str]]
+        Name of columns in DataFrame as partition columns.
+    overwrite: bool
+        Overwrite data if the table / partition already exists.
+    unknown_as_string: bool
+        If True, object type in the DataFrame will be treated as strings.
+        Otherwise errors might be raised.
+    index: bool
+        If True, indexes will be stored. Otherwise they are ignored.
+    index_label: Union[None, str, List[str]]
+        Specify column names for index levels. If absent, level names or default
+        names will be used.
+    lifecycle: Optional[int]
+        Specify lifecycle of the output table.
+
+    Returns
+    -------
+    result: DataFrame
+        Stub DataFrame for execution.
+
+    Notes
+    -----
+    `to_odps_table` returns a stub object for execution. The result returned is
+    not reusable.
+
+    Examples
+    --------
+
+    """
+    if isinstance(table, ODPSTable):
+        table = table.full_table_name
+
+    if isinstance(index_label, str):
+        index_label = [index_label]
+    if isinstance(partition_col, str):
+        partition_col = [partition_col]
+
+    if index_label and len(index_label) != len(df.index.names):
+        raise ValueError(
+            f"index_label needs {len(df.index.nlevels)} labels "
+            f"but it only have {len(index_label)}"
+        )
+
+    op = DataFrameToODPSTable(
+        dtypes=df.dtypes,
+        table_name=table,
+        unknown_as_string=unknown_as_string,
+        partition_spec=partition,
+        partition_columns=partition_col,
+        overwrite=overwrite,
+        index=index,
+        index_label=index_label,
+        lifecycle=lifecycle or options.session.table_lifecycle,
+    )
+    return op(df)

maxframe/dataframe/extensions/__init__.py

@@ -0,0 +1,41 @@
+# 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 ..misc.accessor import CachedAccessor
+from .accessor import (
+    DataFrameMaxFrameAccessor,
+    IndexMaxFrameAccessor,
+    SeriesMaxFrameAccessor,
+)
+from .reshuffle import DataFrameReshuffle, df_reshuffle
+
+
+def _install():
+    from ..core import DATAFRAME_TYPE, INDEX_TYPE, SERIES_TYPE
+
+    DataFrameMaxFrameAccessor._register("reshuffle", df_reshuffle)
+
+    if DataFrameMaxFrameAccessor._api_count:
+        for t in DATAFRAME_TYPE:
+            t.mf = CachedAccessor("mf", DataFrameMaxFrameAccessor)
+    if SeriesMaxFrameAccessor._api_count:
+        for t in SERIES_TYPE:
+            t.mf = CachedAccessor("mf", SeriesMaxFrameAccessor)
+    if IndexMaxFrameAccessor._api_count:
+        for t in INDEX_TYPE:
+            t.mf = CachedAccessor("mf", IndexMaxFrameAccessor)
+
+
+_install()
+del _install

maxframe/dataframe/extensions/accessor.py

@@ -0,0 +1,50 @@
+# 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 functools
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..core import DataFrame, Index, Series
+
+
+class BaseMaxFrameAccessor:
+    def __init__(self, df_obj):
+        self._df_obj = df_obj
+
+    @classmethod
+    def _register(cls, name, func):
+        @functools.wraps(func)
+        def wrapped(self, *args, **kw):
+            return func(self._df_obj, *args, **kw)
+
+        wrapped.__name__ = name
+        setattr(cls, name, wrapped)
+        if hasattr(cls, "_api_count"):  # pragma: no branch
+            cls._api_count += 1
+
+
+class DataFrameMaxFrameAccessor(BaseMaxFrameAccessor):
+    _df_obj: "DataFrame"
+    _api_count: int = 0
+
+
+class SeriesMaxFrameAccessor(BaseMaxFrameAccessor):
+    _df_obj: "Series"
+    _api_count: int = 0
+
+
+class IndexMaxFrameAccessor(BaseMaxFrameAccessor):
+    _df_obj: "Index"
+    _api_count: int = 0

maxframe/dataframe/extensions/reshuffle.py

@@ -0,0 +1,83 @@
+# 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 typing import Any, List, Optional, Union
+
+import pandas as pd
+
+from ... import opcodes
+from ...core import get_output_types
+from ...serialization.serializables import BoolField, ListField
+from ..core import DataFrame, Index, IndexValue, Series
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import parse_index
+
+
+class DataFrameReshuffle(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.DATAFRAME_RESHUFFLE
+
+    group_by = ListField("group_by")
+    sort_by = ListField("sort_by")
+    ascending = BoolField("ascending", default=None)
+    ignore_index = BoolField("ignore_index", default=False)
+
+    def __call__(self, df: Union[DataFrame, Series, Index]):
+        if self.ignore_index:
+            idx_value = parse_index(pd.RangeIndex(0))
+        else:
+            idx_value = df.index_value
+            if isinstance(idx_value.value, IndexValue.RangeIndex):
+                idx_value = parse_index(pd.Int64Index([0]))
+        params = df.params
+        params["index_value"] = idx_value
+        self._output_types = get_output_types(df)
+        return self.new_tileable([df], **params)
+
+
+def df_reshuffle(
+    df_obj,
+    group_by: Optional[List[Any]] = None,
+    sort_by: Optional[List[Any]] = None,
+    ascending: bool = True,
+    ignore_index: bool = False,
+):
+    """
+    Shuffle data in DataFrame or Series to make data distribution more
+    randomized.
+
+    Parameters
+    ----------
+    group_by: Optional[List[Any]]
+        Determine columns to group data while shuffling.
+    sort_by: Optional[List[Any]]
+    ascending
+    ignore_index
+
+    Returns
+    -------
+
+    """
+    if isinstance(group_by, str):
+        group_by = [group_by]
+    if isinstance(sort_by, str):
+        sort_by = [sort_by]
+    if sort_by and not group_by:
+        raise ValueError("to use sort_by requires group_by is specified")
+    op = DataFrameReshuffle(
+        group_by=group_by,
+        sort_by=sort_by,
+        ascending=ascending,
+        ignore_index=ignore_index,
+    )
+    return op(df_obj)

maxframe/dataframe/extensions/tests/__init__.py

@@ -0,0 +1,13 @@
+# 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.

maxframe/dataframe/extensions/tests/test_extensions.py

@@ -0,0 +1,38 @@
+# 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
+import pytest
+
+from .... import dataframe as md
+from ...core import IndexValue
+from ..reshuffle import DataFrameReshuffle
+
+
+def test_reshuffle():
+    df = pd.DataFrame(
+        {"a": [3, 4, 5, 3, 5, 4, 1, 2, 3], "b": [1, 3, 4, 5, 6, 5, 4, 4, 4]}
+    )
+    mdf = md.DataFrame(df, chunk_size=2)
+
+    with pytest.raises(ValueError):
+        mdf.mf.reshuffle(sort_by="b")
+
+    r = mdf.mf.reshuffle()
+    assert isinstance(r.op, DataFrameReshuffle)
+    assert isinstance(r.index_value.value, IndexValue.Int64Index)
+
+    r = mdf.mf.reshuffle(ignore_index=True)
+    assert isinstance(r.op, DataFrameReshuffle)
+    assert isinstance(r.index_value.value, IndexValue.RangeIndex)

maxframe/dataframe/fetch/__init__.py

@@ -0,0 +1,15 @@
+# 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 .core import DataFrameFetch, DataFrameFetchShuffle

maxframe/dataframe/fetch/core.py

@@ -0,0 +1,86 @@
+# 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 ...core import OutputType, register_fetch_class
+from ...core.operator import Fetch, FetchMixin, FetchShuffle
+from ...serialization.serializables import FieldTypes, TupleField
+from ...utils import on_deserialize_shape, on_serialize_shape
+from ..operators import DataFrameOperatorMixin
+
+
+class DataFrameFetchMixin(DataFrameOperatorMixin, FetchMixin):
+    __slots__ = ()
+
+
+class DataFrameFetch(Fetch, DataFrameFetchMixin):
+    # required fields
+    _shape = TupleField(
+        "shape",
+        FieldTypes.tuple,
+        on_serialize=on_serialize_shape,
+        on_deserialize=on_deserialize_shape,
+    )
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+    def _extract_dataframe_or_series_kws(self, kws, **kw):
+        if kws is None:
+            kws = [kw]
+        new_kws = []
+        new_output_types = []
+        for output_type, kwargs in zip(self._output_types, kws):
+            if output_type == OutputType.df_or_series:
+                data_params = kwargs["data_params"]
+                data_type = kwargs["data_type"]
+                if data_type == "series":
+                    new_output_types.append(OutputType.series)
+                else:
+                    new_output_types.append(OutputType.dataframe)
+                new_kws.append(data_params)
+            else:
+                new_output_types.append(output_type)
+                new_kws.append(kwargs)
+        self._output_types = new_output_types
+        return new_kws
+
+    def _new_tileables(self, inputs, kws=None, **kw):
+        if "_key" in kw and self.source_key is None:
+            self.source_key = kw["_key"]
+        new_kws = self._extract_dataframe_or_series_kws(kws, **kw)
+        return super()._new_tileables(inputs, kws=new_kws, **kw)
+
+
+class DataFrameFetchShuffle(FetchShuffle, DataFrameFetchMixin):
+    # required fields
+    _shape = TupleField(
+        "shape",
+        FieldTypes.int64,
+        on_serialize=on_serialize_shape,
+        on_deserialize=on_deserialize_shape,
+    )
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+
+register_fetch_class(OutputType.dataframe, DataFrameFetch, DataFrameFetchShuffle)
+register_fetch_class(
+    OutputType.dataframe_groupby, DataFrameFetch, DataFrameFetchShuffle
+)
+register_fetch_class(OutputType.df_or_series, DataFrameFetch, DataFrameFetchShuffle)
+register_fetch_class(OutputType.series, DataFrameFetch, DataFrameFetchShuffle)
+register_fetch_class(OutputType.series_groupby, DataFrameFetch, DataFrameFetchShuffle)
+register_fetch_class(OutputType.index, DataFrameFetch, DataFrameFetchShuffle)
+register_fetch_class(OutputType.categorical, DataFrameFetch, DataFrameFetchShuffle)

maxframe/dataframe/groupby/__init__.py

@@ -0,0 +1,82 @@
+# 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.
+
+# noinspection PyUnresolvedReferences
+from ..core import DataFrameGroupBy, GroupBy, SeriesGroupBy
+from .core import NamedAgg
+
+
+def _install():
+    from ..core import DATAFRAME_GROUPBY_TYPE, DATAFRAME_TYPE, GROUPBY_TYPE, SERIES_TYPE
+    from .aggregation import agg
+    from .apply import groupby_apply
+    from .core import groupby
+    from .cum import cumcount, cummax, cummin, cumprod, cumsum
+    from .fill import bfill, ffill, fillna
+    from .getitem import df_groupby_getitem
+    from .head import head
+    from .sample import groupby_sample
+    from .transform import groupby_transform
+
+    for cls in DATAFRAME_TYPE:
+        setattr(cls, "groupby", groupby)
+
+    for cls in SERIES_TYPE:
+        setattr(cls, "groupby", groupby)
+
+    for cls in GROUPBY_TYPE:
+        setattr(cls, "agg", agg)
+        setattr(cls, "aggregate", agg)
+
+        setattr(cls, "sum", lambda groupby, **kw: agg(groupby, "sum", **kw))
+        setattr(cls, "prod", lambda groupby, **kw: agg(groupby, "prod", **kw))
+        setattr(cls, "max", lambda groupby, **kw: agg(groupby, "max", **kw))
+        setattr(cls, "min", lambda groupby, **kw: agg(groupby, "min", **kw))
+        setattr(cls, "count", lambda groupby, **kw: agg(groupby, "count", **kw))
+        setattr(cls, "size", lambda groupby, **kw: agg(groupby, "size", **kw))
+        setattr(cls, "mean", lambda groupby, **kw: agg(groupby, "mean", **kw))
+        setattr(cls, "var", lambda groupby, **kw: agg(groupby, "var", **kw))
+        setattr(cls, "std", lambda groupby, **kw: agg(groupby, "std", **kw))
+        setattr(cls, "all", lambda groupby, **kw: agg(groupby, "all", **kw))
+        setattr(cls, "any", lambda groupby, **kw: agg(groupby, "any", **kw))
+        setattr(cls, "skew", lambda groupby, **kw: agg(groupby, "skew", **kw))
+        setattr(cls, "kurt", lambda groupby, **kw: agg(groupby, "kurt", **kw))
+        setattr(cls, "kurtosis", lambda groupby, **kw: agg(groupby, "kurtosis", **kw))
+        setattr(cls, "sem", lambda groupby, **kw: agg(groupby, "sem", **kw))
+        setattr(cls, "nunique", lambda groupby, **kw: agg(groupby, "nunique", **kw))
+
+        setattr(cls, "apply", groupby_apply)
+        setattr(cls, "transform", groupby_transform)
+
+        setattr(cls, "cumcount", cumcount)
+        setattr(cls, "cummin", cummin)
+        setattr(cls, "cummax", cummax)
+        setattr(cls, "cumprod", cumprod)
+        setattr(cls, "cumsum", cumsum)
+
+        setattr(cls, "head", head)
+
+        setattr(cls, "sample", groupby_sample)
+
+        setattr(cls, "ffill", ffill)
+        setattr(cls, "bfill", bfill)
+        setattr(cls, "backfill", bfill)
+        setattr(cls, "fillna", fillna)
+
+    for cls in DATAFRAME_GROUPBY_TYPE:
+        setattr(cls, "__getitem__", df_groupby_getitem)
+
+
+_install()
+del _install