maxframe 1.0.0rc1__cp39-cp39-win_amd64.whl → 1.0.0rc2__cp39-cp39-win_amd64.whl

maxframe/dataframe/indexing/set_index.py

@@ -31,7 +31,7 @@ class DataFrameSetIndex(DataFrameOperator, DataFrameOperatorMixin):
         super().__init__(_output_types=output_types, **kw)
 
     def __call__(self, df):
-        new_df = build_empty_df(df.dtypes).set_index(
+        new_df = build_empty_df(df.dtypes, index=df.index_value.to_pandas()).set_index(
             keys=self.keys,
             drop=self.drop,
             append=self.append,
@@ -47,6 +47,73 @@ class DataFrameSetIndex(DataFrameOperator, DataFrameOperatorMixin):
 
 
 def set_index(df, keys, drop=True, append=False, inplace=False, verify_integrity=False):
+    # TODO add support for set index by series, index, mt.ndarray, etc.
+    """
+    Set the DataFrame index using existing columns.
+
+    Set the DataFrame index (row labels) using one or more existing
+    columns. The index can replace the existing index or expand on it.
+
+    Parameters
+    ----------
+    keys : label or array-like or list of labels
+        This parameter can be either a single column key, or a list containing column keys.
+    drop : bool, default True
+        Delete columns to be used as the new index.
+    append : bool, default False
+        Whether to append columns to existing index.
+    inplace : bool, default False
+        If True, modifies the DataFrame in place (do not create a new object).
+    verify_integrity : bool, default False
+        Check the new index for duplicates. Otherwise defer the check until
+        necessary. Setting to False will improve the performance of this
+        method.
+
+    Returns
+    -------
+    DataFrame or None
+        Changed row labels or None if ``inplace=True``.
+
+    See Also
+    --------
+    DataFrame.reset_index : Opposite of set_index.
+    DataFrame.reindex : Change to new indices or expand indices.
+    DataFrame.reindex_like : Change to same indices as other DataFrame.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+
+    >>> df = md.DataFrame({'month': [1, 4, 7, 10],
+    ...                    'year': [2012, 2014, 2013, 2014],
+    ...                    'sale': [55, 40, 84, 31]})
+    >>> df
+       month  year  sale
+    0      1  2012    55
+    1      4  2014    40
+    2      7  2013    84
+    3     10  2014    31
+
+    Set the index to become the 'month' column:
+
+    >>> df.set_index('month')
+           year  sale
+    month
+    1      2012    55
+    4      2014    40
+    7      2013    84
+    10     2014    31
+
+    Create a MultiIndex using columns 'year' and 'month':
+
+    >>> df.set_index(['year', 'month'])
+                sale
+    year  month
+    2012  1     55
+    2014  4     40
+    2013  7     84
+    2014  10    31
+    """
     op = DataFrameSetIndex(
         keys=keys,
         drop=drop,

maxframe/dataframe/merge/merge.py

@@ -11,12 +11,13 @@
 # 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 abc import abstractmethod
 from collections import namedtuple
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
+from pandas import Index
 
 from ... import opcodes
 from ...core import OutputType
@@ -28,6 +29,7 @@ from ...serialization.serializables import (
     Int32Field,
     KeyField,
     NamedTupleField,
+    Serializable,
     StringField,
     TupleField,
 )
@@ -73,9 +75,208 @@ class DataFrameMergeAlign(MapReduceOperator, DataFrameOperatorMixin):
 MergeSplitInfo = namedtuple("MergeSplitInfo", "split_side, split_index, nsplits")
 
 
+class JoinHint(Serializable):
+    @abstractmethod
+    def verify_params(
+        self,
+        hint_on_df: Union[DataFrame, Series],
+        on: str,
+        is_on_index: bool,
+        how: str,
+        is_hint_for_left: bool,
+    ):
+        pass
+
+    @abstractmethod
+    def verify_can_work_with(self, other: "JoinHint"):
+        pass
+
+
+class MapJoinHint(JoinHint):
+    def verify_params(
+        self,
+        hint_on_df: Union[DataFrame, Series],
+        on: str,
+        is_on_index: bool,
+        how: str,
+        is_hint_for_left: bool,
+    ):
+        if how in ("cross", "outer"):
+            raise ValueError(
+                "Invalid join hint, MapJoinHint is not support in cross and outer join"
+            )
+
+    def verify_can_work_with(self, other: JoinHint):
+        if isinstance(other, SkewJoinHint):
+            raise ValueError(
+                "Invalid join hint, SkewJoinHint cannot work with MapJoinHint"
+            )
+
+
+class DistributedMapJoinHint(JoinHint):
+    shard_count = Int32Field("shard_count")
+    replica_count = Int32Field("replica_count", default=1)
+
+    def verify_params(
+        self,
+        hint_on_df: Union[DataFrame, Series],
+        on: str,
+        is_on_index: bool,
+        how: str,
+        is_hint_for_left: bool,
+    ):
+        if how in ("cross", "outer"):
+            raise ValueError(
+                "Invalid join hint, DistributedMapJoinHint is not support in cross and outer join"
+            )
+        if not hasattr(self, "shard_count"):
+            raise ValueError(
+                "Invalid DistributedMapJoinHint, shard_count must be specified"
+            )
+        if self.shard_count <= 0 or self.replica_count <= 0:
+            raise ValueError(
+                "Invalid DistributedMapJoinHint, shard_count and replica_count must be greater than 0"
+            )
+
+    def verify_can_work_with(self, other: JoinHint):
+        pass
+
+
+class SkewJoinHint(JoinHint):
+    columns = AnyField("columns", default=None)
+
+    @staticmethod
+    def _check_index_levels(index, level_list):
+        selected_levels = set()
+        valid_levels = set(range(index.nlevels))
+        valid_level_names = set(index.names)
+
+        for item in level_list:
+            if isinstance(item, int):
+                if item not in valid_levels:
+                    raise ValueError(f"Level {item} is not a valid index level")
+                if item in selected_levels:
+                    raise ValueError(f"Level {item} is selected multiple times")
+                selected_levels.add(item)
+            elif isinstance(item, str):
+                if item not in valid_level_names:
+                    raise ValueError(f"'{item}' is not a valid index level name")
+                level = index.names.index(item)
+                if level in selected_levels:
+                    raise ValueError(
+                        f"'{item}' (Level {level}) is selected multiple times"
+                    )
+                selected_levels.add(level)
+            else:
+                raise ValueError(f"Invalid input type: {type(item)}")
+
+    @staticmethod
+    def _check_columns(join_on_columns, column_list):
+        selected_columns = set()
+        valid_columns = set(join_on_columns)
+
+        for item in column_list:
+            if isinstance(item, int):
+                if item < 0 or item >= len(join_on_columns):
+                    raise ValueError(f"Column index {item} is out of range")
+                col_name = join_on_columns[item]
+                if col_name in selected_columns:
+                    raise ValueError(
+                        f"Column '{col_name}' (index {item}) is selected multiple times"
+                    )
+                selected_columns.add(col_name)
+            elif isinstance(item, str):
+                if item not in valid_columns:
+                    raise ValueError(f"'{item}' is not a valid column name")
+                if item in selected_columns:
+                    raise ValueError(f"Column '{item}' is selected multiple times")
+                selected_columns.add(item)
+            else:
+                raise ValueError(f"Invalid input type: {type(item)}")
+
+    def verify_params(
+        self,
+        hint_on_df: Union[DataFrame, Series],
+        on: str,
+        is_on_index: bool,
+        how: str,
+        is_hint_for_left: bool,
+    ):
+        if how in ("cross", "outer"):
+            raise ValueError(
+                "Invalid join hint, map join is not support in cross and outer join"
+            )
+        if is_hint_for_left and how == "right":
+            raise ValueError(
+                "Invalid join hint, right join can only use SkewJoinHint on right frame"
+            )
+        elif not is_hint_for_left and how == "left":
+            raise ValueError(
+                "Invalid join hint, left join can only use SkewJoinHint on left frame"
+            )
+
+        # check columns
+        if self.columns is None:
+            return
+
+        if not isinstance(self.columns, list):
+            raise TypeError("Invalid SkewJoinHint, `columns` must be a list")
+
+        if all(isinstance(item, (int, str)) for item in self.columns):
+            # if elements are int (levels) or str (index names or column names)
+            self._verify_valid_index_or_columns(
+                self.columns, hint_on_df.index_value.to_pandas(), on, is_on_index
+            )
+        elif all(isinstance(c, dict) for c in self.columns):
+            # dict with column names and values
+            cols_set = set(self.columns[0].keys())
+            if any(cols_set != set(c.keys()) for c in self.columns):
+                raise ValueError(
+                    "Invalid SkewJoinHint, all values in `columns` need to have same columns"
+                )
+
+            self._verify_valid_index_or_columns(
+                cols_set, hint_on_df.index_value.to_pandas(), on, is_on_index
+            )
+        else:
+            raise TypeError("Invalid SkewJoinHint, annot accept `columns` type")
+
+    def verify_can_work_with(self, other: JoinHint):
+        if isinstance(other, SkewJoinHint):
+            raise ValueError(
+                "Invalid join hint, SkewJoinHint cannot work with MapJoinHint"
+            )
+
+    @staticmethod
+    def _verify_valid_index_or_columns(
+        skew_join_columns: Iterable[Union[int, str]],
+        frame_index: Index,
+        on: Union[str, List[str]],
+        is_on_index: bool,
+    ):
+        if isinstance(on, str):
+            on = [on]
+        on_columns = set(frame_index.names if is_on_index else on)
+        for col in skew_join_columns:
+            if isinstance(col, int):
+                if col < 0 or col >= len(on_columns):
+                    raise ValueError(
+                        f"Invalid, SkeJoinHint, `{col}` is out of join on columns range"
+                    )
+            else:
+                if col not in on_columns:
+                    raise ValueError(
+                        f"Invalid, SkeJoinHint, '{col}' is not a valid column name"
+                    )
+
+
 class DataFrameMerge(DataFrameOperator, DataFrameOperatorMixin):
     _op_type_ = opcodes.DATAFRAME_MERGE
 
+    # workaround for new field since v1.0.0rc2
+    # todo remove this when all versions below v1.0.0rc1 is eliminated
+    _legacy_new_non_primitives = ["left_hint", "right_hint"]
+
     how = StringField("how")
     on = AnyField("on")
     left_on = AnyField("left_on")
@@ -95,6 +296,8 @@ class DataFrameMerge(DataFrameOperator, DataFrameOperatorMixin):
 
     # only for broadcast merge
     split_info = NamedTupleField("split_info")
+    left_hint = AnyField("left_hint", default=None)
+    right_hint = AnyField("right_hint", default=None)
 
     def __init__(self, copy=None, **kwargs):
         super().__init__(copy_=copy, **kwargs)
@@ -165,6 +368,8 @@ def merge(
     auto_merge_threshold: int = 8,
     bloom_filter: Union[bool, str] = "auto",
     bloom_filter_options: Dict[str, Any] = None,
+    left_hint: JoinHint = None,
+    right_hint: JoinHint = None,
 ) -> DataFrame:
     """
     Merge DataFrame or named Series objects with a database-style join.
@@ -267,6 +472,12 @@ def merge(
           when chunk size of left and right is greater than this threshold, apply bloom filter
         * "filter": "large", "small", "both", default "large"
           decides to filter on large, small or both DataFrames.
+    left_hint: JoinHint, default None
+        Join strategy to use for left frame. When data skew occurs, consider these strategies to avoid long-tail issues,
+        but use them cautiously to prevent OOM and unnecessary overhead.
+    right_hint: JoinHint, default None
+        Join strategy to use for right frame.
+
 
     Returns
     -------
@@ -381,6 +592,18 @@ def merge(
                 raise ValueError(
                     f"Invalid filter {k}, available: {BLOOM_FILTER_ON_OPTIONS}"
                 )
+
+    if left_hint:
+        if not isinstance(left_hint, JoinHint):
+            raise TypeError(f"left_hint must be a JoinHint, got {type(left_hint)}")
+        left_hint.verify_can_work_with(right_hint)
+        left_hint.verify_params(df, on or left_on, left_index, how, True)
+
+    if right_hint:
+        if not isinstance(right_hint, JoinHint):
+            raise TypeError(f"right_hint must be a JoinHint, got {type(right_hint)}")
+        right_hint.verify_params(right, on or right_on, right_index, how, False)
+
     op = DataFrameMerge(
         how=how,
         on=on,
@@ -399,6 +622,8 @@ def merge(
         bloom_filter=bloom_filter,
         bloom_filter_options=bloom_filter_options,
         output_types=[OutputType.dataframe],
+        left_hint=left_hint,
+        right_hint=right_hint,
     )
     return op(df, right)
 
@@ -416,6 +641,8 @@ def join(
     auto_merge_threshold: int = 8,
     bloom_filter: Union[bool, Dict] = True,
     bloom_filter_options: Dict[str, Any] = None,
+    left_hint: JoinHint = None,
+    right_hint: JoinHint = None,
 ) -> DataFrame:
     """
     Join columns of another DataFrame.
@@ -480,6 +707,11 @@ def join(
           when chunk size of left and right is greater than this threshold, apply bloom filter
         * "filter": "large", "small", "both", default "large"
           decides to filter on large, small or both DataFrames.
+    left_hint: JoinHint, default None
+        Join strategy to use for left frame. When data skew occurs, consider these strategies to avoid long-tail issues,
+        but use them cautiously to prevent OOM and unnecessary overhead.
+    right_hint: JoinHint, default None
+        Join strategy to use for right frame.
 
     Returns
     -------
@@ -590,4 +822,6 @@ def join(
         auto_merge_threshold=auto_merge_threshold,
         bloom_filter=bloom_filter,
         bloom_filter_options=bloom_filter_options,
+        left_hint=left_hint,
+        right_hint=right_hint,
     )

maxframe/dataframe/merge/tests/test_merge.py

@@ -19,6 +19,7 @@ import pytest
 from ...core import IndexValue
 from ...datasource.dataframe import from_pandas
 from .. import DataFrameMerge, concat
+from ..merge import DistributedMapJoinHint, MapJoinHint, SkewJoinHint
 
 
 def test_merge():
@@ -30,14 +31,39 @@ def test_merge():
     mdf1 = from_pandas(df1, chunk_size=2)
     mdf2 = from_pandas(df2, chunk_size=3)
 
+    mapjoin = MapJoinHint()
+    dist_mapjoin1 = DistributedMapJoinHint(shard_count=5)
+    skew_join1 = SkewJoinHint()
+    skew_join2 = SkewJoinHint(columns=[0])
+    skew_join3 = SkewJoinHint(columns=[{"a": 4}, {"a": 6}])
+    skew_join4 = SkewJoinHint(columns=[{"a": 4, "b": "test"}, {"a": 5, "b": "hello"}])
+
     parameters = [
         {},
         {"how": "left", "right_on": "x", "left_index": True},
+        {
+            "how": "left",
+            "right_on": "x",
+            "left_index": True,
+            "left_hint": mapjoin,
+            "right_hint": mapjoin,
+        },
         {"how": "right", "left_on": "a", "right_index": True},
+        {
+            "how": "right",
+            "left_on": "a",
+            "right_index": True,
+            "left_hint": mapjoin,
+            "right_hint": dist_mapjoin1,
+        },
         {"how": "left", "left_on": "a", "right_on": "x"},
+        {"how": "left", "left_on": "a", "right_on": "x", "left_hint": skew_join1},
         {"how": "right", "left_on": "a", "right_index": True},
+        {"how": "right", "left_on": "a", "right_index": True, "right_hint": skew_join2},
         {"how": "right", "on": "a"},
+        {"how": "right", "on": "a", "right_hint": skew_join3},
         {"how": "inner", "on": ["a", "b"]},
+        {"how": "inner", "on": ["a", "b"], "left_hint": skew_join4},
     ]
 
     for kw in parameters:
@@ -213,3 +239,100 @@ def test_concat():
     mdf2 = from_pandas(df2, chunk_size=3)
     r = concat([mdf1, mdf2], join="inner")
     assert r.shape == (20, 3)
+
+
+def test_invalid_join_hint():
+    df1 = pd.DataFrame(
+        np.arange(20).reshape((4, 5)) + 1, columns=["a", "b", "c", "d", "e"]
+    )
+    df2 = pd.DataFrame(np.arange(20).reshape((5, 4)) + 1, columns=["a", "b", "x", "y"])
+
+    mdf1 = from_pandas(df1, chunk_size=2)
+    mdf2 = from_pandas(df2, chunk_size=3)
+
+    # type error
+    parameters = [
+        {"how": "left", "right_on": "x", "left_index": True, "left_hint": [1]},
+        {
+            "how": "left",
+            "right_on": "x",
+            "left_index": True,
+            "left_hint": {"key": "value"},
+        },
+        {
+            "how": "right",
+            "left_on": "a",
+            "right_index": True,
+            "right_hint": SkewJoinHint(columns=2),
+        },
+        {
+            "how": "left",
+            "left_on": "a",
+            "right_on": "x",
+            "left_hint": SkewJoinHint(columns="a"),
+        },
+        {
+            "how": "right",
+            "left_on": "a",
+            "right_index": True,
+            "right_hint": SkewJoinHint(columns=["0", []]),
+        },
+    ]
+
+    for kw in parameters:
+        print(kw)
+        with pytest.raises(TypeError):
+            mdf1.merge(mdf2, **kw)
+
+    # value error
+    parameters = [
+        # mapjoin can't working with skew join
+        {
+            "how": "left",
+            "right_on": "x",
+            "left_index": True,
+            "left_hint": MapJoinHint(),
+            "right_hint": SkewJoinHint(),
+        },
+        # right join can't apply to skew join left frame
+        {
+            "how": "right",
+            "left_on": "a",
+            "right_index": True,
+            "left_hint": SkewJoinHint(),
+        },
+        # invalid columns
+        {
+            "how": "left",
+            "left_on": "a",
+            "right_on": "x",
+            "left_hint": SkewJoinHint(columns=["b"]),
+        },
+        # invalid index level
+        {
+            "how": "right",
+            "left_on": "a",
+            "right_index": True,
+            "right_hint": SkewJoinHint(columns=[5]),
+        },
+        # unmatched skew join columns
+        {
+            "how": "right",
+            "left_on": "a",
+            "right_index": True,
+            "right_hint": SkewJoinHint(columns=[{0: "value1"}, {1: "value2"}]),
+        },
+        # invalid dist_mapjoin shard_count
+        {"how": "right", "on": "a", "right_hint": DistributedMapJoinHint()},
+        # all can't work with outer join
+        {"how": "outer", "on": ["a", "b"], "left_hint": MapJoinHint()},
+        {
+            "how": "outer",
+            "on": ["a", "b"],
+            "left_hint": DistributedMapJoinHint(shard_count=5),
+        },
+        {"how": "outer", "on": ["a", "b"], "left_hint": SkewJoinHint()},
+    ]
+    for kw in parameters:
+        with pytest.raises(ValueError):
+            mdf1.merge(mdf2, **kw)

maxframe/dataframe/misc/apply.py

@@ -319,6 +319,7 @@ def df_apply(
     skip_infer=False,
     **kwds,
 ):
+    # FIXME: https://github.com/aliyun/alibabacloud-odps-maxframe-client/issues/50
     """
     Apply a function along an axis of the DataFrame.
 
@@ -444,20 +445,12 @@ def df_apply(
     B    27
     dtype: int64
 
-    >>> df.apply(np.sum, axis=1).execute()
+    >>> df.apply(lambda row: int(np.sum(row)), axis=1).execute()
     0    13
     1    13
     2    13
     dtype: int64
 
-    Returning a list-like will result in a Series
-
-    >>> df.apply(lambda x: [1, 2], axis=1).execute()
-    0    [1, 2]
-    1    [1, 2]
-    2    [1, 2]
-    dtype: object
-
     Passing ``result_type='expand'`` will expand list-like results
     to columns of a Dataframe
 
@@ -471,7 +464,7 @@ def df_apply(
     ``result_type='expand'``. The resulting column names
     will be the Series index.
 
-    >>> df.apply(lambda x: md.Series([1, 2], index=['foo', 'bar']), axis=1).execute()
+    >>> df.apply(lambda x: pd.Series([1, 2], index=['foo', 'bar']), axis=1).execute()
        foo  bar
     0    1    2
     1    1    2

maxframe/dataframe/misc/case_when.py

@@ -99,7 +99,7 @@ def case_when(series, caselist):
     >>> b = md.Series([0, 3, 4, 5])
 
     >>> c.case_when(caselist=[(a.gt(0), a),  # condition, replacement
-    ...                       (b.gt(0), b)])
+    ...                       (b.gt(0), b)]).execute()
     0    6
     1    3
     2    1

maxframe/dataframe/misc/describe.py

@@ -15,7 +15,7 @@
 import numpy as np
 import pandas as pd
 
-from ... import opcodes as OperandDef
+from ... import opcodes
 from ...serialization.serializables import AnyField, FieldTypes, KeyField, ListField
 from ..core import SERIES_TYPE
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
@@ -23,7 +23,7 @@ from ..utils import build_empty_df, parse_index
 
 
 class DataFrameDescribe(DataFrameOperator, DataFrameOperatorMixin):
-    _op_type_ = OperandDef.DESCRIBE
+    _op_type_ = opcodes.DESCRIBE
 
     input = KeyField("input", default=None)
     percentiles = ListField("percentiles", FieldTypes.float64, default=None)

maxframe/dataframe/misc/drop_duplicates.py

@@ -37,16 +37,15 @@ class DataFrameDropDuplicates(DuplicateOperand):
             shape += (3,)
         return shape
 
-    @classmethod
-    def _gen_tileable_params(cls, op: "DataFrameDropDuplicates", input_params):
+    def _gen_tileable_params(self, op: "DataFrameDropDuplicates", input_params):
         params = input_params.copy()
-        if op.ignore_index:
+        if op.ignore_index and self._output_types[0] != OutputType.index:
             params["index_value"] = parse_index(pd.RangeIndex(-1))
         else:
             params["index_value"] = gen_unknown_index_value(
                 input_params["index_value"], op.keep, op.subset, type(op).__name__
             )
-        params["shape"] = cls._get_shape(input_params["shape"], op)
+        params["shape"] = self._get_shape(input_params["shape"], op)
         return params
 
     def __call__(self, inp, inplace=False):
@@ -105,6 +104,7 @@ def df_drop_duplicates(
 def series_drop_duplicates(
     series, keep="first", inplace=False, ignore_index=False, method="auto"
 ):
+    # FIXME: https://github.com/aliyun/alibabacloud-odps-maxframe-client/issues/12
     """
     Return Series with duplicate values removed.
 
@@ -148,27 +148,6 @@ def series_drop_duplicates(
     5     hippo
     Name: animal, dtype: object
 
-    With the 'keep' parameter, the selection behaviour of duplicated values
-    can be changed. The value 'first' keeps the first occurrence for each
-    set of duplicated entries. The default value of keep is 'first'.
-
-    >>> s.drop_duplicates().execute()
-    0      lame
-    1       cow
-    3    beetle
-    5     hippo
-    Name: animal, dtype: object
-
-    The value 'last' for parameter 'keep' keeps the last occurrence for
-    each set of duplicated entries.
-
-    >>> s.drop_duplicates(keep='last').execute()
-    1       cow
-    3    beetle
-    4      lame
-    5     hippo
-    Name: animal, dtype: object
-
     The value ``False`` for parameter 'keep' discards all sets of
     duplicated entries. Setting the value of 'inplace' to ``True`` performs
     the operation inplace and returns ``None``.

maxframe/dataframe/misc/eval.py

@@ -120,6 +120,10 @@ class CollectionVisitor(ast.NodeVisitor):
         if obj_name in self.env:
             self.referenced_vars.add(obj_name)
             return self.env[obj_name]
+        try:
+            return self.target[obj_name]
+        except KeyError:
+            pass
         raise KeyError(f"name {obj_name} is not defined")
 
     def visit(self, node):