maxframe 2.0.0b2__cp37-cp37m-win32.whl → 2.3.0rc1__cp37-cp37m-win32.whl

maxframe/dataframe/utils.py

@@ -25,21 +25,22 @@ from typing import TYPE_CHECKING, Any, Callable, List, Optional
 
 import numpy as np
 import pandas as pd
-from pandas.api.types import is_string_dtype
 from pandas.core.dtypes.inference import is_dict_like, is_list_like
 
-from ..core import Entity, ExecutableTuple, OutputType, get_output_types
+from ..config.validators import dtype_backend_validator
+from ..core import ENTITY_TYPE, Entity, ExecutableTuple, OutputType, get_output_types
+from ..lib.dtypes_extension import ExternalBlobDtype, SolidBlob
 from ..lib.mmh3 import hash as mmh_hash
 from ..udf import MarkedFunction
 from ..utils import (
     ModulePlaceholder,
     is_full_slice,
     lazy_import,
-    make_dtype,
     make_dtypes,
     quiet_stdio,
     sbytes,
     tokenize,
+    validate_and_adjust_resource_ratio,
 )
 
 if TYPE_CHECKING:
@@ -57,7 +58,7 @@ cudf = lazy_import("cudf", rename="cudf")
 logger = logging.getLogger(__name__)
 
 try:
-    from pandas import ArrowDtype
+    from ..lib.dtypes_extension import ArrowDtype
 except ImportError:
     ArrowDtype = None
 
@@ -103,9 +104,9 @@ def hash_dtypes(dtypes, size):
     return [dtypes[index] for index in hashed_indexes]
 
 
-def sort_dataframe_inplace(df, *axis):
+def sort_dataframe_inplace(df, *axis, **kw):
     for ax in axis:
-        df.sort_index(axis=ax, inplace=True)
+        df.sort_index(axis=ax, inplace=True, **kw)
     return df
 
 
@@ -456,7 +457,7 @@ def build_split_idx_to_origin_idx(splits, increase=True):
 
 
 def _generate_value(dtype, fill_value):
-    if ArrowDtype and isinstance(dtype, pd.ArrowDtype):
+    if ArrowDtype and isinstance(dtype, ArrowDtype):
         return _generate_value(dtype.pyarrow_dtype, fill_value)
 
     if isinstance(dtype, pa.ListType):
@@ -470,9 +471,19 @@ def _generate_value(dtype, fill_value):
             )
         ]
 
+    if isinstance(dtype, pa.StructType):
+        result = {}
+        for i in range(dtype.num_fields):
+            field = dtype[i]
+            result[field.name] = _generate_value(field.type, fill_value)
+        return result
+
     if isinstance(dtype, pa.DataType):
         return _generate_value(dtype.to_pandas_dtype(), fill_value)
 
+    if isinstance(dtype, ExternalBlobDtype):
+        return SolidBlob(str(fill_value).encode())
+
     # special handle for datetime64 and timedelta64
     dispatch = {
         np.datetime64: pd.Timestamp,
@@ -1012,27 +1023,21 @@ def create_sa_connection(con, **kwargs):
             engine.dispose()
 
 
-def to_arrow_dtypes(dtypes, test_df=None):
-    from .arrays import ArrowStringDtype
+def to_arrow_dtypes(dtypes):
+    from ..io.odpsio.schema import pandas_dtypes_to_arrow_schema
 
+    arrow_schema = pandas_dtypes_to_arrow_schema(dtypes)
     new_dtypes = dtypes.copy()
     for i in range(len(dtypes)):
-        dtype = dtypes.iloc[i]
-        if is_string_dtype(dtype):
-            if test_df is not None:
-                series = test_df.iloc[:, i]
-                # check value
-                non_na_series = series[series.notna()]
-                if len(non_na_series) > 0:
-                    first_value = non_na_series.iloc[0]
-                    if isinstance(first_value, str):
-                        new_dtypes.iloc[i] = ArrowStringDtype()
-                else:  # pragma: no cover
-                    # empty, set arrow string dtype
-                    new_dtypes.iloc[i] = ArrowStringDtype()
-            else:
-                # empty, set arrow string dtype
-                new_dtypes.iloc[i] = ArrowStringDtype()
+        arrow_type = arrow_schema.types[i]
+        dt = dtypes.iloc[i]
+        if isinstance(dt, pd.api.extensions.ExtensionDtype):
+            # make existing extension dtype consistent
+            new_dtypes.iloc[i] = dt
+        elif arrow_type == pa.string():
+            new_dtypes.iloc[i] = pd.StringDtype("pyarrow")
+        else:
+            new_dtypes.iloc[i] = ArrowDtype(arrow_type)
     return new_dtypes
 
 
@@ -1305,7 +1310,7 @@ def pack_func_args(df, funcs, *args, args_bind_position=1, **kwargs) -> Any:
     if is_dict_like(funcs):
         return {k: pack_func_args(df, v, *args, **kwargs) for k, v in funcs.items()}
 
-    if is_list_like(funcs):
+    if is_list_like(funcs) and not isinstance(funcs, ENTITY_TYPE):
         return [pack_func_args(df, v, *args, **kwargs) for v in funcs]
 
     f = get_callable_by_name(df, funcs) if isinstance(funcs, str) else funcs
@@ -1406,23 +1411,54 @@ def infer_dataframe_return_value(
     inherit_index=False,
     build_kw=None,
     elementwise=None,
+    skip_infer=False,
 ) -> InferredDataFrameMeta:
-    from .core import GROUPBY_TYPE
+    from .core import GROUPBY_TYPE, INDEX_TYPE
+    from .typing_ import get_function_output_meta
+
+    unwrapped_func = func
+    if isinstance(unwrapped_func, MarkedFunction):
+        unwrapped_func = unwrapped_func.func
+    while True:
+        if isinstance(unwrapped_func, functools.partial):
+            unwrapped_func = unwrapped_func.func
+        elif hasattr(unwrapped_func, "__wrapped__"):
+            unwrapped_func = unwrapped_func.__wrapped__
+        else:
+            break
+
+    func_annotation_meta = get_function_output_meta(unwrapped_func, df_obj)
+    func_index_value = None
+    if func_annotation_meta:
+        output_type = output_type or func_annotation_meta.output_type
+        dtypes = dtypes if dtypes is not None else func_annotation_meta.dtypes
+        dtype = dtype if dtype is not None else func_annotation_meta.dtype
+        name = name if name is not None else func_annotation_meta.name
+        func_index_value = func_annotation_meta.index_value
+
+    if skip_infer:
+        if isinstance(index, INDEX_TYPE):
+            ret_index_value = index.index_value
+        elif index is not None:
+            ret_index_value = parse_index(index, df_obj.key)
+        else:
+            ret_index_value = func_index_value
+
+        return InferredDataFrameMeta(
+            output_type=output_type,
+            dtypes=dtypes,
+            dtype=dtype,
+            name=name,
+            index_value=ret_index_value,
+        )
+
+    if isinstance(index, INDEX_TYPE):
+        index = index.index_value
 
     if elementwise is None:
-        unwrapped_func = func
-        if isinstance(unwrapped_func, MarkedFunction):
-            unwrapped_func = unwrapped_func.func
-        while True:
-            if isinstance(unwrapped_func, functools.partial):
-                unwrapped_func = unwrapped_func.func
-            elif hasattr(unwrapped_func, "__wrapped__"):
-                unwrapped_func = unwrapped_func.__wrapped__
-            else:
-                break
         elementwise = isinstance(unwrapped_func, np.ufunc)
 
-    ret_index_value = None
+    ret_index_value = func_index_value
     if output_type is not None and (dtypes is not None or dtype is not None):
         if inherit_index:
             ret_index_value = df_obj.index_value
@@ -1439,7 +1475,8 @@ def infer_dataframe_return_value(
                 elementwise=elementwise or False,
             )
 
-    ret_output_type = ret_dtypes = None
+    ret_output_type = None
+    ret_dtypes = dtypes
     maybe_agg = False
     build_kw = build_kw or {}
     obj_key = df_obj.key
@@ -1486,7 +1523,8 @@ def infer_dataframe_return_value(
                     f'please specify `output_type` as "dataframe"'
                 )
             ret_output_type = ret_output_type or OutputType.dataframe
-            ret_dtypes = ret_dtypes or infer_df_obj.dtypes
+            if ret_dtypes is None:
+                ret_dtypes = infer_df_obj.dtypes
         else:
             if output_type is not None and output_type == OutputType.dataframe:
                 raise TypeError(
@@ -1506,7 +1544,7 @@ def infer_dataframe_return_value(
         return InferredDataFrameMeta(
             ret_output_type,
             make_dtypes(ret_dtypes),
-            make_dtype(dtype),
+            make_dtypes(dtype),
             name,
             ret_index_value,
             maybe_agg,
@@ -1519,7 +1557,7 @@ def infer_dataframe_return_value(
         return InferredDataFrameMeta(
             output_type,
             make_dtypes(dtypes),
-            make_dtype(dtype),
+            make_dtypes(dtype),
             name,
             ret_index_value,
             maybe_agg,
@@ -1530,20 +1568,37 @@ def infer_dataframe_return_value(
 def copy_func_scheduling_hints(func, op: "DataFrameOperator") -> None:
     from ..config import options
 
-    if not isinstance(func, MarkedFunction):
-        return
-    if func.expect_engine:
-        op.expect_engine = func.expect_engine
+    expect_engine = None
+    expect_gpu = None
+    default_options = options.function.default_running_options or {}
 
-    expect_resources = func.expect_resources or {}
-    default_function_running_options = options.function.default_running_options or {}
+    if isinstance(func, MarkedFunction):
+        # copy from marked function
+        expect_engine = func.expect_engine
+        expect_resources = func.expect_resources or {}
+        expect_gpu = func.gpu
 
-    for key, value in default_function_running_options.items():
-        if key not in expect_resources or expect_resources.get(key) is None:
-            expect_resources[key] = value
+        # merge default options if not set
+        for key, value in default_options.items():
+            if key not in expect_resources or expect_resources.get(key) is None:
+                expect_resources[key] = value
+    else:
+        # copy from default options
+        expect_resources = default_options
+
+    # Validate and adjust resource ratio constraints on client side
+    expect_resources, _ = validate_and_adjust_resource_ratio(
+        expect_resources,
+        max_memory_cpu_ratio=options.function.allowed_max_memory_cpu_ratio,
+        adjust=True,
+    )
 
-    if func.expect_resources:
+    if expect_engine:
+        op.expect_engine = expect_engine
+    if expect_resources:
         op.expect_resources = expect_resources
+    if expect_gpu:
+        op.gpu = expect_gpu
 
 
 def make_column_list(col, dtypes_or_columns, level=None):
@@ -1576,3 +1631,21 @@ def make_column_list(col, dtypes_or_columns, level=None):
             return idx[mask]
     except (IndexError, TypeError, ValueError):
         return col
+
+
+def call_groupby_with_params(df_or_series, groupby_params: dict):
+    params = groupby_params.copy()
+    selection = params.pop("selection", None)
+    res = df_or_series.groupby(**params)
+    if selection:
+        res = res[selection]
+    return res
+
+
+def validate_dtype_backend(value):
+    if isinstance(value, bool):
+        # compatibility for legacy use_arrow_dtype property
+        value = "pyarrow" if value else "numpy"
+    if not dtype_backend_validator(value):
+        raise ValueError(f"Invalid dtype_backend: {value}")
+    return value

maxframe/dataframe/window/aggregation.py

@@ -19,14 +19,14 @@ import numpy as np
 import pandas as pd
 
 from ...serialization.serializables import AnyField, BoolField, Int32Field, Int64Field
-from ..core import DATAFRAME_TYPE
+from ..core import DATAFRAME_TYPE, ENTITY_TYPE
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import build_df, build_empty_series, parse_index
 
 
 class BaseDataFrameExpandingAgg(DataFrameOperator, DataFrameOperatorMixin):
     min_periods = Int64Field("min_periods", default=None)
-    axis = Int32Field("axis", default=None)
+    axis = Int32Field("axis", default=0)
     func = AnyField("func", default=None)
 
     # always treat count as valid. this behavior is cancelled in pandas 1.0
@@ -52,7 +52,7 @@ class BaseDataFrameExpandingAgg(DataFrameOperator, DataFrameOperatorMixin):
                 index_value = parse_index(
                     test_df.index, expanding.params, inp, store_data=False
                 )
-            self._append_index = test_df.columns.nlevels != empty_df.columns.nlevels
+            self.append_index = test_df.columns.nlevels != empty_df.columns.nlevels
             return self.new_dataframe(
                 [inp],
                 shape=(inp.shape[0], test_df.shape[1]),
@@ -92,5 +92,9 @@ class BaseDataFrameExpandingAgg(DataFrameOperator, DataFrameOperatorMixin):
                 else:
                     new_func[k] = v
             self.func = new_func
-        elif isinstance(self.func, Iterable) and not isinstance(self.func, str):
+        elif (
+            isinstance(self.func, Iterable)
+            and not isinstance(self.func, ENTITY_TYPE)
+            and not isinstance(self.func, str)
+        ):
             self.func = list(self.func)

maxframe/dataframe/window/core.py

@@ -12,11 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...serialization.serializables import KeyField, Serializable
+from ...serialization.serializables import FieldTypes, KeyField, ListField, Serializable
 
 
 class Window(Serializable):
+    _mf_specific_fields = ["order_cols", "ascending"]
+
     input = KeyField("input", default=None)
+    order_cols = ListField("order_cols", default=None)
+    ascending = ListField("ascending", FieldTypes.bool, default=None)
+
+    def __init__(self, *, order_cols=None, ascending=True, **kwargs):
+        if order_cols and not isinstance(order_cols, list):
+            order_cols = [order_cols]
+        if not isinstance(ascending, list):
+            ascending = [ascending]
+        elif order_cols and len(order_cols) != len(ascending):
+            raise ValueError("order_cols and ascending must have same length")
+        super().__init__(order_cols=order_cols, ascending=ascending, **kwargs)
 
     @property
     def params(self):

maxframe/dataframe/window/ewm.py

@@ -36,8 +36,6 @@ _window_has_method = pd_release_version >= (1, 4, 0)
 class DataFrameEwmAgg(BaseDataFrameExpandingAgg):
     _op_type_ = opcodes.EWM_AGG
 
-    _exec_cache = dict()
-
     alpha = Float64Field("alpha")
     adjust = BoolField("adjust")
     alpha_ignore_na = BoolField("alpha_ignore_na")
@@ -234,7 +232,7 @@ def ewm(
         raise ValueError("alpha must satisfy: 0 < alpha <= 1")
 
     if alpha == 1:
-        return obj.expanding(min_periods=min_periods, axis=axis)
+        return obj.expanding(min_periods=min_periods)
 
     if _default_min_period_1:
         min_periods = min_periods or 1

maxframe/dataframe/window/expanding.py

@@ -12,49 +12,53 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 from collections import OrderedDict
 
 from ... import opcodes
-from ...serialization.serializables import (
-    BoolField,
-    Int32Field,
-    Int64Field,
-    StringField,
-)
-from ...utils import pd_release_version
-from ..utils import validate_axis
+from ...serialization.serializables import BoolField, Int64Field
 from .aggregation import BaseDataFrameExpandingAgg
 from .core import Window
 
-_window_has_method = pd_release_version >= (1, 3, 0)
-_window_has_center = pd_release_version < (2, 0, 0)
-
 
 class DataFrameExpandingAgg(BaseDataFrameExpandingAgg):
     _op_type_ = opcodes.EXPANDING_AGG
 
-    center = BoolField("center", default=None)
+    def __init__(self, *args, **kw):
+        # suspend MF-specific args by now
+        for key in Expanding._mf_specific_fields:
+            kw.pop(key, None)
+        super().__init__(*args, **kw)
 
 
 class Expanding(Window):
+    _mf_specific_fields = Window._mf_specific_fields + ["shift", "reverse_range"]
+
     min_periods = Int64Field("min_periods")
-    axis = Int32Field("axis")
-    center = BoolField("center")
-    method = StringField("method", default="single")
+    # MF specific argument for position shift of window
+    shift = Int64Field("shift", default=None)
+    # MF specific argument for reversed window (sort of "narrowing")
+    reverse_range = BoolField("reverse_range", default=False)
 
     def __call__(self, df):
-        return df.expanding(**self.params)
+        try:
+            return df.expanding(**self.params)
+        except TypeError:
+            params = (self.params or dict()).copy()
+            for key in self._mf_specific_fields:
+                params.pop(key, None)
+            return df.expanding(**params)
 
     @property
     def params(self):
         p = OrderedDict()
 
-        args = ["min_periods", "center", "axis", "method"]
-        if not _window_has_method:  # pragma: no cover
-            args = [a for a in args if a != "method"]
-        if not _window_has_center:
-            args = [a for a in args if a != "center"]
+        args = [
+            "min_periods",
+            "shift",
+            "reverse_range",
+            "order_cols",
+            "ascending",
+        ]
 
         for k in args:
             p[k] = getattr(self, k)
@@ -73,6 +77,9 @@ class Expanding(Window):
     def sum(self):
         return self.aggregate("sum")
 
+    def prod(self):
+        return self.aggregate("prod")
+
     def count(self):
         return self.aggregate("count")
 
@@ -85,14 +92,14 @@ class Expanding(Window):
     def mean(self):
         return self.aggregate("mean")
 
-    def var(self):
-        return self.aggregate("var")
+    def var(self, **kwargs):
+        return self.aggregate("var", **kwargs)
 
-    def std(self):
-        return self.aggregate("std")
+    def std(self, **kwargs):
+        return self.aggregate("std", **kwargs)
 
 
-def expanding(obj, min_periods=1, center=False, axis=0):
+def expanding(obj, min_periods=1, shift=0, reverse_range=False):
     """
     Provide expanding transformations.
 
@@ -139,11 +146,6 @@ def expanding(obj, min_periods=1, center=False, axis=0):
     3  3.0
     4  7.0
     """
-    axis = validate_axis(axis, obj)
-
-    if center:
-        raise NotImplementedError("center == True is not supported")
-    if axis == 1:
-        raise NotImplementedError("axis other than 0 is not supported")
-
-    return Expanding(input=obj, min_periods=min_periods, center=center, axis=axis)
+    return Expanding(
+        input=obj, min_periods=min_periods, shift=shift, reverse_range=reverse_range
+    )

maxframe/dataframe/window/rolling.py

@@ -23,38 +23,42 @@ from ...serialization.serializables import (
     AnyField,
     BoolField,
     DictField,
+    FieldTypes,
     Int32Field,
     Int64Field,
     KeyField,
+    ListField,
     StringField,
     TupleField,
 )
-from ...utils import pd_release_version
 from ..core import DATAFRAME_TYPE
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import build_empty_df, build_empty_series, parse_index, validate_axis
 from .core import Window
 
-_window_has_method = pd_release_version >= (1, 3, 0)
-_with_pandas_issue_38908 = pd_release_version == (1, 2, 0)
-
 
 class DataFrameRollingAgg(DataFrameOperator, DataFrameOperatorMixin):
     _op_type_ = opcodes.ROLLING_AGG
 
     input = KeyField("input")
-    window = AnyField("window")
-    min_periods = Int64Field("min_periods")
-    center = BoolField("center")
-    win_type = StringField("win_type")
-    on = StringField("on")
-    axis = Int32Field("axis")
-    closed = StringField("closed")
-    func = AnyField("func")
-    func_args = TupleField("func_args")
-    func_kwargs = DictField("func_kwargs")
+    window = AnyField("window", default=None)
+    min_periods = Int64Field("min_periods", default=None)
+    center = BoolField("center", default=None)
+    win_type = StringField("win_type", default=None)
+    on = StringField("on", default=None)
+    axis = Int32Field("axis", default=None)
+    closed = StringField("closed", default=None)
+    func = AnyField("func", default=None)
+    func_args = TupleField("func_args", default=None)
+    func_kwargs = DictField("func_kwargs", default=None)
+    # for chunks
+    preds = ListField("preds", FieldTypes.key, default=None)
+    succs = ListField("succs", FieldTypes.key, default=None)
 
     def __init__(self, output_types=None, **kw):
+        # suspend MF-specific args by now
+        for key in Rolling._mf_specific_fields:
+            kw.pop(key, None)
         super().__init__(_output_types=output_types, **kw)
 
     @classmethod
@@ -62,6 +66,10 @@ class DataFrameRollingAgg(DataFrameOperator, DataFrameOperatorMixin):
         super()._set_inputs(op, inputs)
         input_iter = iter(op._inputs)
         op.input = next(input_iter)
+        if op.preds is not None:
+            op.preds = [next(input_iter) for _ in op.preds]
+        if op.succs is not None:
+            op.succs = [next(input_iter) for _ in op.succs]
 
     def __call__(self, rolling):
         inp = rolling.input
@@ -74,6 +82,8 @@ class DataFrameRollingAgg(DataFrameOperator, DataFrameOperatorMixin):
                 params["win_type"] = None
             if self.func != "count":
                 empty_df = empty_df._get_numeric_data()
+            for key in Rolling._mf_specific_fields:
+                params.pop(key, None)
             test_df = empty_df.rolling(**params).agg(self.func)
             if self.axis == 0:
                 index_value = inp.index_value
@@ -93,7 +103,10 @@ class DataFrameRollingAgg(DataFrameOperator, DataFrameOperatorMixin):
             empty_series = build_empty_series(
                 inp.dtype, index=pd_index[:0], name=inp.name
             )
-            test_obj = empty_series.rolling(**rolling.params).agg(self.func)
+            rolling_params = rolling.params.copy()
+            for k in Rolling._mf_specific_fields:
+                rolling_params.pop(k, None)
+            test_obj = empty_series.rolling(**rolling_params).agg(self.func)
             if isinstance(test_obj, pd.DataFrame):
                 return self.new_dataframe(
                     [inp],
@@ -113,6 +126,8 @@ class DataFrameRollingAgg(DataFrameOperator, DataFrameOperatorMixin):
 
 
 class Rolling(Window):
+    _mf_specific_fields = Window._mf_specific_fields + ["shift"]
+
     window = AnyField("window", default=None)
     min_periods = Int64Field("min_periods", default=None)
     center = BoolField("center", default=None)
@@ -120,33 +135,25 @@ class Rolling(Window):
     on = StringField("on", default=None)
     axis = Int32Field("axis", default=None)
     closed = StringField("closed", default=None)
-    method = StringField("method", default="single")
+    # MF specific argument for position shift of window
+    shift = Int64Field("shift", default=None)
 
     @property
     def params(self):
         p = OrderedDict()
 
-        if not _window_has_method:  # pragma: no cover
-            args = [
-                "window",
-                "min_periods",
-                "center",
-                "win_type",
-                "axis",
-                "on",
-                "closed",
-            ]
-        else:
-            args = [
-                "window",
-                "min_periods",
-                "center",
-                "win_type",
-                "axis",
-                "on",
-                "closed",
-                "method",
-            ]
+        args = [
+            "window",
+            "min_periods",
+            "center",
+            "win_type",
+            "axis",
+            "on",
+            "closed",
+            "shift",
+            "order_cols",
+            "ascending",
+        ]
 
         for attr in args:
             p[attr] = getattr(self, attr)
@@ -164,8 +171,11 @@ class Rolling(Window):
             empty_obj = build_empty_series(
                 self.input.dtype, index=pd_index[:0], name=self.input.name
             )
-        pd_rolling = empty_obj.rolling(**self.params)
-        for k in self.params:
+        params = (self.params or dict()).copy()
+        for key in self._mf_specific_fields:
+            params.pop(key, None)
+        pd_rolling = empty_obj.rolling(**params)
+        for k in params:
             # update value according to pandas rolling
             setattr(self, k, getattr(pd_rolling, k))
 

maxframe/dataframe/window/tests/test_expanding.py

@@ -23,15 +23,9 @@ def test_expanding():
     df = pd.DataFrame(np.random.rand(4, 3), columns=list("abc"))
     df2 = md.DataFrame(df)
 
-    with pytest.raises(NotImplementedError):
-        _ = df2.expanding(3, center=True)
-
-    with pytest.raises(NotImplementedError):
-        _ = df2.expanding(3, axis=1)
-
     r = df2.expanding(3)
     expected = df.expanding(3)
-    assert repr(r) == repr(expected)
+    assert repr(r).split(",", 1)[0] == repr(expected).split(",", 1)[0]
 
     assert "b" in dir(r)