maxframe 0.1.0b5__cp38-cp38-macosx_10_9_universal2.whl → 1.0.0rc2__cp38-cp38-macosx_10_9_universal2.whl

maxframe/dataframe/datastore/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

maxframe/dataframe/datastore/tests/test_to_odps.py

@@ -0,0 +1,48 @@
+# 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 pytest
+
+from ... import DataFrame
+from ..to_odps import to_odps_table
+
+
+@pytest.fixture
+def df():
+    return DataFrame({"A": [1, 2], "B": [3, 4]})
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"partition_col": ["A", "C"]},
+        {"partition_col": "C"},
+        {"partition": "a=1,C=2"},
+    ],
+)
+def test_to_odps_table_validation(df, kwargs):
+    with pytest.raises(ValueError):
+        to_odps_table(df, "test_table", **kwargs)
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"partition_col": ["a", "B"]},
+        {"partition_col": "a"},
+        {"partition": "C=1,d=2"},
+    ],
+)
+def test_to_odps_table_vaild(df, kwargs):
+    to_odps_table(df, "test_table", **kwargs)

maxframe/dataframe/datastore/to_odps.py

@@ -18,10 +18,12 @@ import logging
 from typing import List, Optional, Union
 
 from odps.models import Table as ODPSTable
+from odps.types import PartitionSpec
 
 from ... import opcodes
 from ...config import options
 from ...core import OutputType
+from ...odpsio import build_dataframe_table_meta
 from ...serialization.serializables import (
     BoolField,
     FieldTypes,
@@ -147,6 +149,25 @@ def to_odps_table(
             f"index_label needs {len(df.index.nlevels)} labels "
             f"but it only have {len(index_label)}"
         )
+    table_cols = set(build_dataframe_table_meta(df).table_column_names)
+    if partition:
+        partition_intersect = (
+            set(x.lower() for x in PartitionSpec(partition).keys()) & table_cols
+        )
+        if partition_intersect:
+            raise ValueError(
+                f"Data column(s) {partition_intersect} in the dataframe"
+                " cannot be used in parameter 'partition'."
+                " Use 'partition_col' instead."
+            )
+
+    if partition_col:
+        partition_diff = set(x.lower() for x in partition_col) - table_cols
+        if partition_diff:
+            raise ValueError(
+                f"Partition column(s) {partition_diff}"
+                " is not the data column(s) of the input dataframe."
+            )
 
     op = DataFrameToODPSTable(
         dtypes=df.dtypes,

maxframe/dataframe/groupby/cum.py

@@ -59,7 +59,6 @@ class GroupByCumReductionOperator(DataFrameOperatorMixin, DataFrameOperator):
         out_dtypes = self._calc_out_dtypes(groupby)
 
         kw = in_df.params.copy()
-        kw["index_value"] = parse_index(pd.RangeIndex(-1), groupby.key)
         if self.output_types[0] == OutputType.dataframe:
             kw.update(
                 dict(

maxframe/dataframe/groupby/tests/test_groupby.py

@@ -282,14 +282,17 @@ def test_groupby_cum():
         r = getattr(mdf.groupby("b"), fun)()
         assert r.op.output_types[0] == OutputType.dataframe
         assert r.shape == (len(df1), 2)
+        assert r.index_value.key == mdf.index_value.key
 
         r = getattr(mdf.groupby("b"), fun)(axis=1)
         assert r.op.output_types[0] == OutputType.dataframe
         assert r.shape == (len(df1), 3)
+        assert r.index_value.key == mdf.index_value.key
 
     r = mdf.groupby("b").cumcount()
     assert r.op.output_types[0] == OutputType.series
     assert r.shape == (len(df1),)
+    assert r.index_value.key == mdf.index_value.key
 
     series1 = pd.Series([2, 2, 5, 7, 3, 7, 8, 8, 5, 6])
     ms1 = md.Series(series1, chunk_size=3)
@@ -298,6 +301,7 @@ def test_groupby_cum():
         r = getattr(ms1.groupby(lambda x: x % 2), fun)()
         assert r.op.output_types[0] == OutputType.series
         assert r.shape == (len(series1),)
+        assert r.index_value.key == ms1.index_value.key
 
 
 def test_groupby_fill():

maxframe/dataframe/indexing/add_prefix_suffix.py

@@ -51,7 +51,7 @@ def _get_prefix_suffix_docs(is_prefix: bool):
     Examples
     --------
     >>> import maxframe.dataframe as md
-    >>> s = md.Series([1, 2, 3, 4])
+        >>> s = md.Series([1, 2, 3, 4])
     >>> s.execute()
     0    1
     1    2

maxframe/dataframe/indexing/align.py

@@ -138,7 +138,7 @@ class DataFrameAlign(DataFrameOperator, DataFrameOperatorMixin):
                 series_index = rhs.index_value.to_pandas()
                 dtypes = lhs.dtypes.reindex(
                     lhs.dtypes.index.join(series_index, how=self.join)
-                ).fillna(np.dtype(np.float_))
+                ).fillna(np.dtype(float))
                 l_shape[1] = r_size = len(dtypes)
                 col_val = r_idx_val = parse_index(dtypes.index, store_data=True)
 

maxframe/dataframe/indexing/rename.py

@@ -17,7 +17,7 @@ import warnings
 from ... import opcodes
 from ...core import get_output_types
 from ...serialization.serializables import AnyField, StringField
-from ..core import SERIES_TYPE
+from ..core import INDEX_TYPE, SERIES_TYPE
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import build_df, build_series, parse_index, validate_axis
 
@@ -73,6 +73,8 @@ class DataFrameRename(DataFrameOperator, DataFrameOperatorMixin):
             params["index_value"] = parse_index(new_index)
         if df.ndim == 1:
             params["name"] = new_df.name
+            if isinstance(df, INDEX_TYPE):
+                params["names"] = new_df.names
         return self.new_tileable([df], **params)
 
 
@@ -303,11 +305,6 @@ def series_rename(
     1    2
     2    3
     Name: my_name, dtype: int64
-    >>> s.rename(lambda x: x ** 2).execute()  # function, changes labels.execute()
-    0    1
-    1    2
-    4    3
-    dtype: int64
     >>> s.rename({1: 3, 2: 5}).execute()  # mapping, changes labels.execute()
     0    1
     3    2
@@ -410,37 +407,6 @@ def index_set_names(index, names, level=None, inplace=False):
     See Also
     --------
     Index.rename : Able to set new names without level.
-
-    Examples
-    --------
-    >>> import maxframe.dataframe as md
-    >>> idx = md.Index([1, 2, 3, 4])
-    >>> idx.execute()
-    Int64Index([1, 2, 3, 4], dtype='int64')
-    >>> idx.set_names('quarter').execute()
-    Int64Index([1, 2, 3, 4], dtype='int64', name='quarter')
-
-    >>> idx = md.MultiIndex.from_product([['python', 'cobra'],
-    ...                                   [2018, 2019]])
-    >>> idx.execute()
-    MultiIndex([('python', 2018),
-                ('python', 2019),
-                ( 'cobra', 2018),
-                ( 'cobra', 2019)],
-               )
-    >>> idx.set_names(['kind', 'year'], inplace=True)
-    >>> idx.execute()
-    MultiIndex([('python', 2018),
-                ('python', 2019),
-                ( 'cobra', 2018),
-                ( 'cobra', 2019)],
-               names=['kind', 'year'])
-    >>> idx.set_names('species', level=0).execute()
-    MultiIndex([('python', 2018),
-                ('python', 2019),
-                ( 'cobra', 2018),
-                ( 'cobra', 2019)],
-               names=['species', 'year'])
     """
     op = DataFrameRename(
         index_mapper=names, level=level, output_types=get_output_types(index)

maxframe/dataframe/indexing/sample.py

@@ -195,7 +195,6 @@ def sample(
             num_legs  num_wings  num_specimen_seen
     falcon         2          2                 10
     fish           0          0                  8
-
     """
     axis = validate_axis(axis or 0, df_or_series)
     if axis == 1:

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,
     )