maxframe 0.1.0b4__cp39-cp39-win32.whl → 1.0.0rc1__cp39-cp39-win32.whl

maxframe/dataframe/misc/case_when.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 numpy as np
+from pandas.core.dtypes.cast import find_common_type
+
+from ... import opcodes
+from ...core import TILEABLE_TYPE
+from ...serialization.serializables import FieldTypes, ListField
+from ..core import SERIES_TYPE
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import apply_if_callable
+
+
+class DataFrameCaseWhen(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.CASE_WHEN
+
+    conditions = ListField("conditions", FieldTypes.reference, default=None)
+    replacements = ListField("replacements", FieldTypes.reference, default=None)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        it = iter(inputs)
+        next(it)
+        self.conditions = [
+            next(it) if isinstance(t, TILEABLE_TYPE) else t for t in self.conditions
+        ]
+        self.replacements = [
+            next(it) if isinstance(t, TILEABLE_TYPE) else t for t in self.replacements
+        ]
+
+    def __call__(self, series):
+        replacement_dtypes = [
+            it.dtype if isinstance(it, SERIES_TYPE) else np.array(it).dtype
+            for it in self.replacements
+        ]
+        dtype = find_common_type([series.dtype] + replacement_dtypes)
+
+        condition_tileables = [
+            it for it in self.conditions if isinstance(it, TILEABLE_TYPE)
+        ]
+        replacement_tileables = [
+            it for it in self.replacements if isinstance(it, TILEABLE_TYPE)
+        ]
+        inputs = [series] + condition_tileables + replacement_tileables
+
+        params = series.params
+        params["dtype"] = dtype
+        return self.new_series(inputs, **params)
+
+
+def case_when(series, caselist):
+    """
+    Replace values where the conditions are True.
+
+    Parameters
+    ----------
+    caselist : A list of tuples of conditions and expected replacements
+        Takes the form:  ``(condition0, replacement0)``,
+        ``(condition1, replacement1)``, ... .
+        ``condition`` should be a 1-D boolean array-like object
+        or a callable. If ``condition`` is a callable,
+        it is computed on the Series
+        and should return a boolean Series or array.
+        The callable must not change the input Series
+        (though pandas doesn`t check it). ``replacement`` should be a
+        1-D array-like object, a scalar or a callable.
+        If ``replacement`` is a callable, it is computed on the Series
+        and should return a scalar or Series. The callable
+        must not change the input Series.
+
+    Returns
+    -------
+    Series
+
+    See Also
+    --------
+    Series.mask : Replace values where the condition is True.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> c = md.Series([6, 7, 8, 9], name='c')
+    >>> a = md.Series([0, 0, 1, 2])
+    >>> b = md.Series([0, 3, 4, 5])
+
+    >>> c.case_when(caselist=[(a.gt(0), a),  # condition, replacement
+    ...                       (b.gt(0), b)])
+    0    6
+    1    3
+    2    1
+    3    2
+    Name: c, dtype: int64
+    """
+    if not isinstance(caselist, list):
+        raise TypeError(
+            f"The caselist argument should be a list; instead got {type(caselist)}"
+        )
+
+    if not caselist:
+        raise ValueError(
+            "provide at least one boolean condition, "
+            "with a corresponding replacement."
+        )
+
+    for num, entry in enumerate(caselist):
+        if not isinstance(entry, tuple):
+            raise TypeError(
+                f"Argument {num} must be a tuple; instead got {type(entry)}."
+            )
+        if len(entry) != 2:
+            raise ValueError(
+                f"Argument {num} must have length 2; "
+                "a condition and replacement; "
+                f"instead got length {len(entry)}."
+            )
+    caselist = [
+        (
+            apply_if_callable(condition, series),
+            apply_if_callable(replacement, series),
+        )
+        for condition, replacement in caselist
+    ]
+    conditions = [case[0] for case in caselist]
+    replacements = [case[1] for case in caselist]
+    op = DataFrameCaseWhen(conditions=conditions, replacements=replacements)
+    return op(series)

maxframe/dataframe/misc/memory_usage.py

@@ -58,7 +58,7 @@ class DataFrameMemoryUsage(DataFrameOperator, DataFrameOperatorMixin):
         """
         if df_or_series.ndim == 1:
             # the input data is a series, a Scalar will be returned
-            return self.new_scalar([df_or_series], dtype=np.dtype(np.int_))
+            return self.new_scalar([df_or_series], dtype=np.dtype(int))
         else:
             # the input data is a DataFrame, a Scalar will be returned
             # calculate shape of returning series given ``op.index``
@@ -71,7 +71,7 @@ class DataFrameMemoryUsage(DataFrameOperator, DataFrameOperatorMixin):
                 [df_or_series],
                 index_value=self._adapt_index(df_or_series.columns_value),
                 shape=new_shape,
-                dtype=np.dtype(np.int_),
+                dtype=np.dtype(int),
             )
 
 

maxframe/dataframe/misc/pivot_table.py

@@ -0,0 +1,262 @@
+# 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 pandas.api.types import is_list_like
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import AnyField, BoolField, StringField
+from ...utils import no_default
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import build_df, parse_index
+
+
+class DataFramePivotTable(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.PIVOT_TABLE
+
+    values = AnyField("values", default=None)
+    index = AnyField("index", default=None)
+    columns = AnyField("columns", default=None)
+    aggfunc = AnyField("aggfunc", default="mean")
+    fill_value = AnyField("fill_value", default=None)
+    margins = BoolField("margins", default=False)
+    dropna = BoolField("dropna", default=True)
+    margins_name = StringField("margins_name", default=None)
+    sort = BoolField("sort", default=False)
+
+    def __init__(self, **kw):
+        super().__init__(**kw)
+        self.output_types = [OutputType.dataframe]
+
+    def __call__(self, df):
+        index_value = columns_value = dtypes = None
+        if self.index is not None:
+            # index is now a required field
+            if len(self.index) == 1:
+                index_data = pd.Index(
+                    [], dtype=df.dtypes[self.index[0]], name=self.index[0]
+                )
+            else:
+                index_data = pd.MultiIndex.from_frame(build_df(df[self.index]))
+            index_value = parse_index(index_data)
+
+            if self.columns is None:  # output columns can be determined
+                sel_df = df
+                groupby_obj = sel_df.groupby(self.index)
+                if self.values:
+                    groupby_obj = groupby_obj[self.values]
+                aggregated_df = groupby_obj.agg(self.aggfunc)
+                index_value = aggregated_df.index_value
+                columns_value = aggregated_df.columns_value
+                dtypes = aggregated_df.dtypes
+            else:
+                columns_value = dtypes = None
+        return self.new_dataframe(
+            [df],
+            shape=(np.nan, np.nan),
+            dtypes=dtypes,
+            columns_value=columns_value,
+            index_value=index_value,
+        )
+
+
+def pivot_table(
+    data,
+    values=None,
+    index=None,
+    columns=None,
+    aggfunc="mean",
+    fill_value=None,
+    margins=False,
+    dropna=True,
+    margins_name="All",
+    sort=True,
+):
+    """
+    Create a spreadsheet-style pivot table as a DataFrame.
+
+    The levels in the pivot table will be stored in MultiIndex objects
+    (hierarchical indexes) on the index and columns of the result DataFrame.
+
+    Parameters
+    ----------
+    values : column to aggregate, optional
+    index : column, Grouper, array, or list of the previous
+        If an array is passed, it must be the same length as the data. The
+        list can contain any of the other types (except list).
+        Keys to group by on the pivot table index.  If an array is passed,
+        it is being used as the same manner as column values.
+    columns : column, Grouper, array, or list of the previous
+        If an array is passed, it must be the same length as the data. The
+        list can contain any of the other types (except list).
+        Keys to group by on the pivot table column.  If an array is passed,
+        it is being used as the same manner as column values.
+    aggfunc : function, list of functions, dict, default numpy.mean
+        If list of functions passed, the resulting pivot table will have
+        hierarchical columns whose top level are the function names
+        (inferred from the function objects themselves)
+        If dict is passed, the key is column to aggregate and value
+        is function or list of functions.
+    fill_value : scalar, default None
+        Value to replace missing values with (in the resulting pivot table,
+        after aggregation).
+    margins : bool, default False
+        Add all row / columns (e.g. for subtotal / grand totals).
+    dropna : bool, default True
+        Do not include columns whose entries are all NaN.
+    margins_name : str, default 'All'
+        Name of the row / column that will contain the totals
+        when margins is True.
+    sort : bool, default True
+        Specifies if the result should be sorted.
+
+    Returns
+    -------
+    DataFrame
+        An Excel style pivot table.
+
+    See Also
+    --------
+    DataFrame.pivot : Pivot without aggregation that can handle
+        non-numeric data.
+    DataFrame.melt: Unpivot a DataFrame from wide to long format,
+        optionally leaving identifiers set.
+    wide_to_long : Wide panel to long format. Less flexible but more
+        user-friendly than melt.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({"A": ["foo", "foo", "foo", "foo", "foo",
+    ...                          "bar", "bar", "bar", "bar"],
+    ...                    "B": ["one", "one", "one", "two", "two",
+    ...                          "one", "one", "two", "two"],
+    ...                    "C": ["small", "large", "large", "small",
+    ...                          "small", "large", "small", "small",
+    ...                          "large"],
+    ...                    "D": [1, 2, 2, 3, 3, 4, 5, 6, 7],
+    ...                    "E": [2, 4, 5, 5, 6, 6, 8, 9, 9]})
+    >>> df.execute()
+         A    B      C  D  E
+    0  foo  one  small  1  2
+    1  foo  one  large  2  4
+    2  foo  one  large  2  5
+    3  foo  two  small  3  5
+    4  foo  two  small  3  6
+    5  bar  one  large  4  6
+    6  bar  one  small  5  8
+    7  bar  two  small  6  9
+    8  bar  two  large  7  9
+
+    This first example aggregates values by taking the sum.
+
+    >>> table = md.pivot_table(df, values='D', index=['A', 'B'],
+    ...                        columns=['C'], aggfunc=np.sum)
+    >>> table.execute()
+    C        large  small
+    A   B
+    bar one    4.0    5.0
+        two    7.0    6.0
+    foo one    4.0    1.0
+        two    NaN    6.0
+
+    We can also fill missing values using the `fill_value` parameter.
+
+    >>> table = md.pivot_table(df, values='D', index=['A', 'B'],
+    ...                        columns=['C'], aggfunc=np.sum, fill_value=0)
+    >>> table.execute()
+    C        large  small
+    A   B
+    bar one      4      5
+        two      7      6
+    foo one      4      1
+        two      0      6
+
+    The next example aggregates by taking the mean across multiple columns.
+
+    >>> table = md.pivot_table(df, values=['D', 'E'], index=['A', 'C'],
+    ...                        aggfunc={'D': np.mean,
+    ...                                 'E': np.mean})
+    >>> table.execute()
+                    D         E
+    A   C
+    bar large  5.500000  7.500000
+        small  5.500000  8.500000
+    foo large  2.000000  4.500000
+        small  2.333333  4.333333
+
+    We can also calculate multiple types of aggregations for any given
+    value column.
+
+    >>> table = md.pivot_table(df, values=['D', 'E'], index=['A', 'C'],
+    ...                        aggfunc={'D': np.mean,
+    ...                                 'E': [min, max, np.mean]})
+    >>> table.execute()
+                    D    E
+                mean  max      mean  min
+    A   C
+    bar large  5.500000  9.0  7.500000  6.0
+        small  5.500000  9.0  8.500000  8.0
+    foo large  2.000000  5.0  4.500000  4.0
+        small  2.333333  6.0  4.333333  2.0
+    """
+    if index is None and columns is None:
+        raise ValueError(
+            "No group keys passed, need to specify at least one of index or columns"
+        )
+
+    def make_col_list(col):
+        try:
+            if col in data.dtypes.index:
+                return [col]
+        except TypeError:
+            return col
+        return col
+
+    values_list = make_col_list(values)
+    index_list = make_col_list(index)
+    columns_list = make_col_list(columns)
+
+    name_to_attr = {"values": values_list, "index": index_list, "columns": columns_list}
+    for key, val in name_to_attr.items():
+        if val is None:
+            continue
+        if not is_list_like(val):
+            raise ValueError(f"Need to specify {key} as a list-like object.")
+        non_exist_key = next((c for c in val if c not in data.dtypes.index), no_default)
+        if non_exist_key is not no_default:
+            raise ValueError(
+                f"Column {non_exist_key} specified in {key} is not a valid column."
+            )
+
+    if columns is None and not margins:
+        if values_list:
+            data = data[index_list + values_list]
+        return data.groupby(index, sort=sort).agg(aggfunc)
+
+    op = DataFramePivotTable(
+        values=values,
+        index=index,
+        columns=columns,
+        aggfunc=aggfunc,
+        fill_value=fill_value,
+        margins=margins,
+        dropna=dropna,
+        margins_name=margins_name,
+        sort=sort,
+    )
+    return op(data)

maxframe/dataframe/misc/tests/test_misc.py

@@ -18,9 +18,11 @@ import pytest
 
 from .... import opcodes
 from ....core import OutputType
+from ....dataframe import DataFrame
 from ....tensor.core import TENSOR_TYPE
 from ... import eval as maxframe_eval
 from ... import get_dummies, to_numeric
+from ...arithmetic import DataFrameGreater, DataFrameLess
 from ...core import CATEGORICAL_TYPE, DATAFRAME_TYPE, INDEX_TYPE, SERIES_TYPE
 from ...datasource.dataframe import from_pandas as from_pandas_df
 from ...datasource.index import from_pandas as from_pandas_index
@@ -405,3 +407,85 @@ def test_to_numeric():
 
     with pytest.raises(ValueError):
         _ = to_numeric([])
+
+
+def test_case_when():
+    rs = np.random.RandomState(0)
+    raw = pd.DataFrame(
+        rs.randint(1000, size=(20, 8)), columns=["c" + str(i + 1) for i in range(8)]
+    )
+    df = from_pandas_df(raw, chunk_size=8)
+
+    with pytest.raises(TypeError):
+        df.c1.case_when(df.c2)
+    with pytest.raises(ValueError):
+        df.c1.case_when([])
+    with pytest.raises(TypeError):
+        df.c1.case_when([[]])
+    with pytest.raises(ValueError):
+        df.c1.case_when([()])
+
+    col = df.c1.case_when([(df.c2 < 10, 10), (df.c2 > 20, df.c3)])
+    assert len(col.inputs) == 4
+    assert isinstance(col.inputs[1].op, DataFrameLess)
+    assert isinstance(col.inputs[2].op, DataFrameGreater)
+
+
+def test_apply():
+    df = DataFrame({"a": [1, 2, 3], "b": [1, 2, 3], "c": [1, 2, 3]})
+
+    keys = [1, 2]
+
+    def f(x, keys):
+        if x["a"] in keys:
+            return [1, 0]
+        else:
+            return [0, 1]
+
+    apply_df = df[["a"]].apply(
+        f,
+        output_type="dataframe",
+        dtypes=pd.Series(["int64", "int64"]),
+        axis=1,
+        result_type="expand",
+        keys=keys,
+    )
+    assert apply_df.shape == (3, 2)
+
+
+def test_pivot_table():
+    from ...groupby.aggregation import DataFrameGroupByAgg
+    from ...misc.pivot_table import DataFramePivotTable
+
+    raw = pd.DataFrame(
+        {
+            "A": "foo foo foo foo foo bar bar bar bar".split(),
+            "B": "one one one two two one one two two".split(),
+            "C": "small large large small small large small small large".split(),
+            "D": [1, 2, 2, 3, 3, 4, 5, 6, 7],
+            "E": [2, 4, 5, 5, 6, 6, 8, 9, 9],
+        }
+    )
+    df = from_pandas_df(raw, chunk_size=8)
+    with pytest.raises(ValueError):
+        df.pivot_table(index=123)
+    with pytest.raises(ValueError):
+        df.pivot_table(index=["F"])
+    with pytest.raises(ValueError):
+        df.pivot_table(values=["D", "E"], aggfunc="sum")
+
+    t = df.pivot_table(index="A")
+    assert isinstance(t.op, DataFrameGroupByAgg)
+    t = df.pivot_table(index="A", values=["D", "E"], aggfunc="sum")
+    assert isinstance(t.op, DataFrameGroupByAgg)
+
+    t = df.pivot_table(index=["A", "B"], values=["D", "E"], aggfunc="sum", margins=True)
+    assert isinstance(t.op, DataFramePivotTable)
+
+    t = df.pivot_table(index="A", columns=["B", "C"], aggfunc="sum")
+    assert isinstance(t.op, DataFramePivotTable)
+    assert t.shape == (np.nan, np.nan)
+
+    t = df.pivot_table(index=["A", "B"], columns="C", aggfunc="sum")
+    assert isinstance(t.op, DataFramePivotTable)
+    assert t.shape == (np.nan, np.nan)

maxframe/dataframe/plotting/core.py

@@ -17,7 +17,7 @@ from collections import OrderedDict
 import pandas as pd
 
 from ...core import ENTITY_TYPE, ExecutableTuple
-from ...utils import adapt_docstring
+from ...utils import adapt_docstring, get_item_if_scalar
 
 
 class PlotAccessor:
@@ -34,7 +34,7 @@ class PlotAccessor:
             .fetch(session=session)
         )
         for p, v in zip(to_executes, executed):
-            result[p] = v
+            result[p] = get_item_if_scalar(v)
 
         data = result.pop("__object__")
         pd_kwargs = kwargs.copy()

maxframe/dataframe/reduction/core.py

@@ -30,7 +30,7 @@ from ...serialization.serializables import (
     StringField,
 )
 from ...typing_ import TileableType
-from ...utils import pd_release_version, tokenize
+from ...utils import get_item_if_scalar, pd_release_version, tokenize
 from ..operators import DATAFRAME_TYPE, DataFrameOperator, DataFrameOperatorMixin
 from ..utils import (
     build_df,
@@ -715,6 +715,7 @@ class ReductionCompiler:
                 keys_to_vars = {inp.key: local_key_to_var[inp.key] for inp in t.inputs}
 
                 def _interpret_var(v):
+                    v = get_item_if_scalar(v)
                     # get representation for variables
                     if hasattr(v, "key"):
                         return keys_to_vars[v.key]

maxframe/dataframe/statistics/corr.py

@@ -43,7 +43,7 @@ class DataFrameCorr(DataFrameOperator, DataFrameOperatorMixin):
     def __call__(self, df_or_series):
         if isinstance(df_or_series, SERIES_TYPE):
             inputs = filter_inputs([df_or_series, self.other])
-            return self.new_scalar(inputs, dtype=np.dtype(np.float_))
+            return self.new_scalar(inputs, dtype=np.dtype(float))
         else:
 
             def _filter_numeric(obj):
@@ -60,7 +60,7 @@ class DataFrameCorr(DataFrameOperator, DataFrameOperatorMixin):
             inputs = filter_inputs([df_or_series, self.other])
             if self.axis is None:
                 dtypes = pd.Series(
-                    [np.dtype(np.float_)] * len(df_or_series.dtypes),
+                    [np.dtype(float)] * len(df_or_series.dtypes),
                     index=df_or_series.dtypes.index,
                 )
                 return self.new_dataframe(
@@ -85,7 +85,7 @@ class DataFrameCorr(DataFrameOperator, DataFrameOperatorMixin):
                 return self.new_series(
                     inputs,
                     shape=shape,
-                    dtype=np.dtype(np.float_),
+                    dtype=np.dtype(float),
                     index_value=new_index_value,
                 )
 

maxframe/dataframe/utils.py

@@ -1136,6 +1136,13 @@ def concat_on_columns(objs: List) -> Any:
     return result
 
 
+def apply_if_callable(maybe_callable, obj, **kwargs):
+    if callable(maybe_callable):
+        return maybe_callable(obj, **kwargs)
+
+    return maybe_callable
+
+
 def patch_sa_engine_execute():
     """
     pandas did not resolve compatibility issue of sqlalchemy 2.0, the issue

maxframe/errors.py

@@ -17,5 +17,18 @@ class MaxFrameError(Exception):
     pass
 
 
+class MaxFrameIntentionalError(MaxFrameError):
+    pass
+
+
 class MaxFrameUserError(MaxFrameError):
     pass
+
+
+class NoTaskServerResponseError(MaxFrameError):
+    pass
+
+
+class SessionAlreadyClosedError(MaxFrameError):
+    def __init__(self, session_id: str):
+        super().__init__(f"Session {session_id} is already closed")

maxframe/extension.py

@@ -48,6 +48,18 @@ class MaxFrameExtension(metaclass=abc.ABCMeta):
         """
         pass
 
+    @classmethod
+    async def reload_session(cls, session_id: str) -> None:
+        """
+        Reload the session state when the session is recovered from failover.
+
+        Parameters
+        ----------
+        session_id : str
+            The session id.
+        """
+        pass
+
     @classmethod
     def init_service_extension(cls) -> None:
         """

maxframe/learn/contrib/utils.py

@@ -0,0 +1,52 @@
+# 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 sys
+
+
+def make_import_error_func(package_name):
+    def _func(*_, **__):  # pragma: no cover
+        raise ImportError(
+            f"Cannot import {package_name}, please reinstall that package."
+        )
+
+    return _func
+
+
+def config_mod_getattr(mod_dict, globals_):
+    def __getattr__(name):
+        import importlib
+
+        if name in mod_dict:
+            mod_name, cls_name = mod_dict[name].rsplit(".", 1)
+            mod = importlib.import_module(mod_name, globals_["__name__"])
+            cls = globals_[name] = getattr(mod, cls_name)
+            return cls
+        else:  # pragma: no cover
+            raise AttributeError(name)
+
+    if sys.version_info[:2] < (3, 7):
+        for _mod in mod_dict.keys():
+            __getattr__(_mod)
+
+    def __dir__():
+        return sorted([n for n in globals_ if not n.startswith("_")] + list(mod_dict))
+
+    globals_.update(
+        {
+            "__getattr__": __getattr__,
+            "__dir__": __dir__,
+            "__all__": list(__dir__()),
+            "__warningregistry__": dict(),
+        }
+    )