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

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)
@@ -150,7 +353,7 @@ def merge(
     df: Union[DataFrame, Series],
     right: Union[DataFrame, Series],
     how: str = "inner",
-    on: str = None,
+    on: Union[str, List[str]] = None,
     left_on: str = None,
     right_on: str = None,
     left_index: bool = False,
@@ -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

@@ -16,9 +16,11 @@ import numpy as np
 import pandas as pd
 import pytest
 
+from ....tests.utils import assert_mf_index_dtype
 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 +32,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:
@@ -135,7 +162,7 @@ def test_append():
     adf = mdf1.append(mdf2)
 
     assert adf.shape == (20, 4)
-    assert isinstance(adf.index_value.value, IndexValue.Int64Index)
+    assert_mf_index_dtype(adf.index_value.value, np.int64)
 
     mdf1 = from_pandas(df1, chunk_size=3)
     mdf2 = from_pandas(df2, chunk_size=3)
@@ -155,6 +182,7 @@ def test_concat():
     r = concat([mdf1, mdf2], axis="index")
 
     assert r.shape == (20, 4)
+    assert not isinstance(r.index_value.to_pandas(), pd.RangeIndex)
     pd.testing.assert_series_equal(r.dtypes, df1.dtypes)
 
     df3 = pd.DataFrame(
@@ -213,3 +241,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/__init__.py

@@ -14,6 +14,7 @@
 
 from .apply import df_apply, series_apply
 from .astype import astype, index_astype
+from .case_when import case_when
 from .check_monotonic import (
     check_monotonic,
     is_monotonic,
@@ -37,6 +38,7 @@ from .map import index_map, series_map
 from .melt import melt
 from .memory_usage import df_memory_usage, index_memory_usage, series_memory_usage
 from .pct_change import pct_change
+from .pivot_table import pivot_table
 from .qcut import qcut
 from .select_dtypes import select_dtypes
 from .shift import shift, tshift
@@ -69,6 +71,7 @@ def _install():
         setattr(t, "melt", melt)
         setattr(t, "memory_usage", df_memory_usage)
         setattr(t, "pct_change", pct_change)
+        setattr(t, "pivot_table", pivot_table)
         setattr(t, "pop", df_pop)
         setattr(t, "query", df_query)
         setattr(t, "select_dtypes", select_dtypes)
@@ -81,6 +84,7 @@ def _install():
     for t in SERIES_TYPE:
         setattr(t, "apply", series_apply)
         setattr(t, "astype", astype)
+        setattr(t, "case_when", case_when)
         setattr(t, "check_monotonic", check_monotonic)
         setattr(t, "describe", describe)
         setattr(t, "diff", series_diff)

maxframe/dataframe/misc/apply.py

@@ -170,6 +170,8 @@ class ApplyOperator(
         elif self.output_types[0] == OutputType.dataframe:
             shape = [np.nan, np.nan]
             shape[1 - self.axis] = df.shape[1 - self.axis]
+            if self.axis == 1:
+                shape[1] = len(dtypes)
             shape = tuple(shape)
         else:
             shape = (df.shape[1 - self.axis],)
@@ -225,7 +227,7 @@ class ApplyOperator(
                 else:  # pragma: no cover
                     index_value = parse_index(infer_series.index)
             else:
-                index_value = parse_index(None, series)
+                index_value = parse_index(series.index_value)
 
             if output_type == OutputType.dataframe:
                 if dtypes is None:
@@ -317,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.
 
@@ -442,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
 
@@ -469,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