maxframe 2.0.0b2__cp311-cp311-win_amd64.whl → 2.2.0__cp311-cp311-win_amd64.whl

maxframe/dataframe/reduction/unique.py

@@ -28,12 +28,19 @@ from ...tensor.core import TensorOrder
 from ...utils import lazy_import
 from ..core import DATAFRAME_TYPE
 from ..initializer import Series as asseries
-from .core import CustomReduction, DataFrameReductionMixin, DataFrameReductionOperator
+from .core import (
+    CustomReduction,
+    DataFrameReduction,
+    DataFrameReductionMixin,
+    ReductionCallable,
+)
 
 cudf = lazy_import("cudf")
 
 
 class UniqueReduction(CustomReduction):
+    _func_name = "unique"
+
     def agg(self, data):  # noqa: W0221  # pylint: disable=arguments-differ
         xdf = cudf if self.is_gpu() else pd
         # convert to series data
@@ -43,7 +50,12 @@ class UniqueReduction(CustomReduction):
         return data.unique()
 
 
-class DataFrameUnique(DataFrameReductionOperator, DataFrameReductionMixin):
+class UniqueReductionCallable(ReductionCallable):
+    def __call__(self, value):
+        return UniqueReduction(name="unique", is_gpu=self.kwargs["is_gpu"])(value)
+
+
+class DataFrameUnique(DataFrameReduction, DataFrameReductionMixin):
     _op_type_ = opcodes.UNIQUE
     _func_name = "unique"
 
@@ -53,9 +65,14 @@ class DataFrameUnique(DataFrameReductionOperator, DataFrameReductionMixin):
     def is_atomic(self):
         return True
 
+    def get_reduction_args(self, axis=None):
+        return {}
+
     @classmethod
     def get_reduction_callable(cls, op):
-        return UniqueReduction(name=cls._func_name, is_gpu=op.is_gpu())
+        return UniqueReductionCallable(
+            func_name=cls._func_name, kwargs=dict(is_gpu=op.is_gpu())
+        )
 
     def __call__(self, a):
         if not isinstance(a, ENTITY_TYPE):

maxframe/dataframe/reduction/var.py

@@ -15,10 +15,21 @@
 from ... import opcodes
 from ...core import OutputType
 from ...serialization.serializables import Int32Field
-from .core import DataFrameReductionMixin, DataFrameReductionOperator
+from .core import DataFrameReduction, DataFrameReductionMixin, ReductionCallable
 
 
-class DataFrameVar(DataFrameReductionOperator, DataFrameReductionMixin):
+class VarReductionCallable(ReductionCallable):
+    def __call__(self, value):
+        skipna, ddof = self.kwargs["skipna"], self.kwargs["ddof"]
+        cnt = value.count()
+        if ddof == 0:
+            return (value**2).mean(skipna=skipna) - (value.mean(skipna=skipna)) ** 2
+        return (
+            (value**2).sum(skipna=skipna) - value.sum(skipna=skipna) ** 2 / cnt
+        ) / (cnt - ddof)
+
+
+class DataFrameVar(DataFrameReduction, DataFrameReductionMixin):
     _op_type_ = opcodes.VAR
     _func_name = "var"
 
@@ -27,16 +38,9 @@ class DataFrameVar(DataFrameReductionOperator, DataFrameReductionMixin):
     @classmethod
     def get_reduction_callable(cls, op: "DataFrameVar"):
         skipna, ddof = op.skipna, op.ddof
-
-        def var(x):
-            cnt = x.count()
-            if ddof == 0:
-                return (x**2).mean(skipna=skipna) - (x.mean(skipna=skipna)) ** 2
-            return ((x**2).sum(skipna=skipna) - x.sum(skipna=skipna) ** 2 / cnt) / (
-                cnt - ddof
-            )
-
-        return var
+        return VarReductionCallable(
+            func_name="var", kwargs={"skipna": skipna, "ddof": ddof}
+        )
 
 
 def var_series(series, axis=None, skipna=True, level=None, ddof=1, method=None):

maxframe/dataframe/reshape/__init__.py

@@ -0,0 +1,38 @@
+# Copyright 1999-2025 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.
+
+from .melt import melt
+from .pivot import pivot
+from .pivot_table import pivot_table
+from .stack import stack
+from .unstack import unstack
+
+
+def _install():
+    from ..core import DATAFRAME_TYPE, SERIES_TYPE
+
+    for t in DATAFRAME_TYPE:
+        setattr(t, "melt", melt)
+        setattr(t, "pivot", pivot)
+        setattr(t, "pivot_table", pivot_table)
+        setattr(t, "stack", stack)
+        setattr(t, "unstack", unstack)
+
+    for t in SERIES_TYPE:
+        setattr(t, "unstack", unstack)
+
+
+_install()
+del _install

maxframe/dataframe/{misc → reshape}/pivot.py

@@ -35,6 +35,7 @@ class DataFramePivot(DataFrameOperator, DataFrameOperatorMixin):
     columns = AnyField("columns", default=None)
 
     agg_results = KeyField("agg_results", default=None)
+    fill_value = AnyField("fill_value", default=None)
 
     def __init__(self, aggfunc=None, **kw):
         if aggfunc is None:

maxframe/dataframe/{misc → reshape}/pivot_table.py

@@ -39,6 +39,7 @@ class DataFramePivotTable(DataFrameOperator, DataFrameOperatorMixin):
     margins_name = StringField("margins_name", default=None)
     sort = BoolField("sort", default=False)
 
+    # generate intermediate agg result to infer value of columns
     agg_results = KeyField("agg_results", default=None)
 
     def __init__(self, aggfunc=None, **kw):

maxframe/dataframe/reshape/unstack.py

@@ -0,0 +1,114 @@
+# Copyright 1999-2025 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.
+
+from ...udf import builtin_function
+from ..core import SERIES_TYPE
+
+
+@builtin_function
+def _adjust_names(pivoted, idx_names=None, col_names=None, remove_col_level=False):
+    if remove_col_level:
+        pivoted = pivoted.droplevel(0, axis=1)
+    if idx_names:
+        pivoted = pivoted.rename_axis(idx_names, axis=0)
+    if col_names:
+        pivoted = pivoted.rename_axis(col_names, axis=1)
+    return pivoted
+
+
+def unstack(df_or_series, level=-1, fill_value=None):
+    """
+    Unstack, also known as pivot, Series with MultiIndex to produce DataFrame.
+
+    Parameters
+    ----------
+    level : int, str, or list of these, default last level
+        Level(s) to unstack, can pass level name.
+    fill_value : scalar value, default None
+        Value to use when replacing NaN values.
+
+    Returns
+    -------
+    DataFrame
+        Unstacked Series.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> s = md.Series([1, 2, 3, 4],
+    ...               index=md.MultiIndex.from_product([['one', 'two'],
+    ...                                                 ['a', 'b']]))
+    >>> s.execute()
+    one  a    1
+         b    2
+    two  a    3
+         b    4
+    dtype: int64
+
+    >>> s.unstack(level=-1).execute()
+         a  b
+    one  1  2
+    two  3  4
+
+    >>> s.unstack(level=0).execute()
+         one  two
+    a    1    3
+    b    2    4
+    """
+    if df_or_series.index.nlevels == 1:
+        if isinstance(df_or_series, SERIES_TYPE):
+            raise ValueError("index must be a MultiIndex to unstack")
+        raw_idx_nlevels = df_or_series.index.nlevels
+        stacked = df_or_series.stack(level=level, fill_value=fill_value)
+        orig_order = list(stacked.index.nlevels)
+        new_order = orig_order[raw_idx_nlevels:] + orig_order[:raw_idx_nlevels]
+        return stacked.reorder_levels(new_order)
+
+    nlevels = df_or_series.index.nlevels
+    idx_names = list(df_or_series.index.names)
+    if df_or_series.ndim == 1:
+        col_names = []
+    else:
+        col_names = list(df_or_series.columns.names)
+
+    name_to_idx = {name: i for i, name in enumerate(idx_names)}
+    level = level if isinstance(level, list) else [level]
+    level_idxes = [(nlevels + name_to_idx.get(lv, lv)) % nlevels for lv in level]
+    level_idx_set = set(level_idxes)
+    level_cols = [f"__level_{idx}" for idx in range(nlevels)]
+
+    if df_or_series.ndim == 1:
+        data = df_or_series.to_frame(name="__data")
+    else:
+        data = df_or_series
+    pivot_val_col = list(data.dtypes.index)
+    pivot_col_col = [level_cols[i] for i in level_idxes]
+    pivot_idx_col = [level_cols[i] for i in range(nlevels) if i not in level_idx_set]
+
+    new_idx_names = [idx_names[i] for i in range(nlevels) if i not in level_idx_set]
+    new_col_names = col_names + [idx_names[i] for i in level_idxes]
+
+    data = data.reset_index(names=level_cols)
+    pivoted = data.pivot(
+        index=pivot_idx_col, columns=pivot_col_col, values=pivot_val_col
+    )
+    pivoted.op.fill_value = fill_value
+    return pivoted.mf.apply_chunk(
+        _adjust_names,
+        col_names=new_col_names,
+        idx_names=new_idx_names,
+        remove_col_level=df_or_series.ndim == 1,
+        skip_infer=True,
+        output_type="dataframe",
+    )

maxframe/dataframe/sort/__init__.py

@@ -18,14 +18,22 @@ from .sort_values import DataFrameSortValues
 
 def _install():
     from ..core import DATAFRAME_TYPE, SERIES_TYPE
+    from .argsort import series_argsort
+    from .nlargest import df_nlargest, series_nlargest
+    from .nsmallest import df_nsmallest, series_nsmallest
     from .sort_index import sort_index
     from .sort_values import dataframe_sort_values, series_sort_values
 
     for cls in DATAFRAME_TYPE:
+        setattr(cls, "nlargest", df_nlargest)
+        setattr(cls, "nsmallest", df_nsmallest)
         setattr(cls, "sort_values", dataframe_sort_values)
         setattr(cls, "sort_index", sort_index)
 
     for cls in SERIES_TYPE:
+        setattr(cls, "argsort", series_argsort)
+        setattr(cls, "nlargest", series_nlargest)
+        setattr(cls, "nsmallest", series_nsmallest)
         setattr(cls, "sort_values", series_sort_values)
         setattr(cls, "sort_index", sort_index)
 

maxframe/dataframe/sort/argsort.py

@@ -0,0 +1,62 @@
+# Copyright 1999-2025 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.
+
+
+def series_argsort(series, axis=0, kind="quicksort", order=None, stable=None):
+    """
+    Return the integer indices that would sort the Series values.
+
+    Override ndarray.argsort. Argsorts the value, omitting NA/null values,
+    and places the result in the same locations as the non-NA values.
+
+    Parameters
+    ----------
+    axis : {0 or 'index'}
+        Unused. Parameter needed for compatibility with DataFrame.
+    kind : {'mergesort', 'quicksort', 'heapsort', 'stable'}, default 'quicksort'
+        Choice of sorting algorithm. See :func:`numpy.sort` for more
+        information. 'mergesort' and 'stable' are the only stable algorithms.
+    order : None
+        Has no effect but is accepted for compatibility with numpy.
+    stable : None
+        Has no effect but is accepted for compatibility with numpy.
+
+    Returns
+    -------
+    Series[np.intp]
+        Positions of values within the sort order with -1 indicating
+        nan values.
+
+    See Also
+    --------
+    numpy.ndarray.argsort : Returns the indices that would sort this array.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> import maxframe.dataframe as md
+    >>> s = md.Series([3, 2, 1])
+    >>> s.argsort().execute()
+    0    2
+    1    1
+    2    0
+    dtype: int64
+    """
+    from ... import tensor as mt
+    from ..datasource.from_tensor import series_from_tensor
+
+    _ = axis, order, stable
+    axis = 0
+    t = mt.argsort(series.to_tensor(), axis=axis, kind=kind)
+    return series_from_tensor(t, index=series.index)

maxframe/dataframe/sort/core.py

@@ -34,3 +34,4 @@ class DataFrameSortOperator(DataFrameOperator):
     parallel_kind = StringField("parallel_kind")
     psrs_kinds = ListField("psrs_kinds", FieldTypes.string)
     nrows = Int64Field("nrows", default=None)
+    keep_kind = StringField("keep_kind", default="head")

maxframe/dataframe/sort/nlargest.py

@@ -0,0 +1,238 @@
+# Copyright 1999-2025 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.
+
+from ...core import OutputType
+from .sort_values import DataFrameSortValues
+
+
+def _nlargest(df, n, columns=None, keep="first"):
+    op = DataFrameSortValues(
+        output_types=[OutputType.dataframe],
+        axis=0,
+        by=columns,
+        ignore_index=False,
+        ascending=False,
+        nrows=n,
+        keep_kind=keep,
+    )
+    return op(df)
+
+
+def df_nlargest(df, n, columns, keep="first"):
+    """
+    Return the first `n` rows ordered by `columns` in descending order.
+
+    Return the first `n` rows with the largest values in `columns`, in
+    descending order. The columns that are not specified are returned as
+    well, but not used for ordering.
+
+    This method is equivalent to
+    ``df.sort_values(columns, ascending=False).head(n)``, but more
+    performant.
+
+    Parameters
+    ----------
+    n : int
+        Number of rows to return.
+    columns : label or list of labels
+        Column label(s) to order by.
+    keep : {'first', 'last', 'all'}, default 'first'
+        Where there are duplicate values:
+
+        - `first` : prioritize the first occurrence(s)
+        - `last` : prioritize the last occurrence(s)
+        - ``all`` : do not drop any duplicates, even it means
+                    selecting more than `n` items.
+
+    Returns
+    -------
+    DataFrame
+        The first `n` rows ordered by the given columns in descending
+        order.
+
+    See Also
+    --------
+    DataFrame.nsmallest : Return the first `n` rows ordered by `columns` in
+        ascending order.
+    DataFrame.sort_values : Sort DataFrame by the values.
+    DataFrame.head : Return the first `n` rows without re-ordering.
+
+    Notes
+    -----
+    This function cannot be used with all column types. For example, when
+    specifying columns with `object` or `category` dtypes, ``TypeError`` is
+    raised.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({'population': [59000000, 65000000, 434000,
+    ...                                   434000, 434000, 337000, 11300,
+    ...                                   11300, 11300],
+    ...                    'GDP': [1937894, 2583560 , 12011, 4520, 12128,
+    ...                            17036, 182, 38, 311],
+    ...                    'alpha-2': ["IT", "FR", "MT", "MV", "BN",
+    ...                                "IS", "NR", "TV", "AI"]},
+    ...                   index=["Italy", "France", "Malta",
+    ...                          "Maldives", "Brunei", "Iceland",
+    ...                          "Nauru", "Tuvalu", "Anguilla"])
+    >>> df.execute()
+              population      GDP alpha-2
+    Italy       59000000  1937894      IT
+    France      65000000  2583560      FR
+    Malta         434000    12011      MT
+    Maldives      434000     4520      MV
+    Brunei        434000    12128      BN
+    Iceland       337000    17036      IS
+    Nauru          11300      182      NR
+    Tuvalu         11300       38      TV
+    Anguilla       11300      311      AI
+
+    In the following example, we will use ``nlargest`` to select the three
+    rows having the largest values in column "population".
+
+    >>> df.nlargest(3, 'population').execute()
+            population      GDP alpha-2
+    France    65000000  2583560      FR
+    Italy     59000000  1937894      IT
+    Malta       434000    12011      MT
+
+    When using ``keep='last'``, ties are resolved in reverse order:
+
+    >>> df.nlargest(3, 'population', keep='last').execute()
+            population      GDP alpha-2
+    France    65000000  2583560      FR
+    Italy     59000000  1937894      IT
+    Brunei      434000    12128      BN
+
+    When using ``keep='all'``, all duplicate items are maintained:
+
+    >>> df.nlargest(3, 'population', keep='all').execute()
+              population      GDP alpha-2
+    France      65000000  2583560      FR
+    Italy       59000000  1937894      IT
+    Malta         434000    12011      MT
+    Maldives      434000     4520      MV
+    Brunei        434000    12128      BN
+
+    To order by the largest values in column "population" and then "GDP",
+    we can specify multiple columns like in the next example.
+
+    >>> df.nlargest(3, ['population', 'GDP']).execute()
+            population      GDP alpha-2
+    France    65000000  2583560      FR
+    Italy     59000000  1937894      IT
+    Brunei      434000    12128      BN
+    """
+    return _nlargest(df, n, columns, keep=keep)
+
+
+def series_nlargest(df, n, keep="first"):
+    """
+    Return the largest `n` elements.
+
+    Parameters
+    ----------
+    n : int, default 5
+        Return this many descending sorted values.
+    keep : {'first', 'last', 'all'}, default 'first'
+        When there are duplicate values that cannot all fit in a
+        Series of `n` elements:
+
+        - ``first`` : return the first `n` occurrences in order
+            of appearance.
+        - ``last`` : return the last `n` occurrences in reverse
+            order of appearance.
+        - ``all`` : keep all occurrences. This can result in a Series of
+            size larger than `n`.
+
+    Returns
+    -------
+    Series
+        The `n` largest values in the Series, sorted in decreasing order.
+
+    See Also
+    --------
+    Series.nsmallest: Get the `n` smallest elements.
+    Series.sort_values: Sort Series by values.
+    Series.head: Return the first `n` rows.
+
+    Notes
+    -----
+    Faster than ``.sort_values(ascending=False).head(n)`` for small `n`
+    relative to the size of the ``Series`` object.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> countries_population = {"Italy": 59000000, "France": 65000000,
+    ...                         "Malta": 434000, "Maldives": 434000,
+    ...                         "Brunei": 434000, "Iceland": 337000,
+    ...                         "Nauru": 11300, "Tuvalu": 11300,
+    ...                         "Anguilla": 11300, "Montserrat": 5200}
+    >>> s = md.Series(countries_population)
+    >>> s.execute()
+    Italy       59000000
+    France      65000000
+    Malta         434000
+    Maldives      434000
+    Brunei        434000
+    Iceland       337000
+    Nauru          11300
+    Tuvalu         11300
+    Anguilla       11300
+    Montserrat      5200
+    dtype: int64
+
+    The `n` largest elements where ``n=5`` by default.
+
+    >>> s.nlargest().execute()
+    France      65000000
+    Italy       59000000
+    Malta         434000
+    Maldives      434000
+    Brunei        434000
+    dtype: int64
+
+    The `n` largest elements where ``n=3``. Default `keep` value is 'first'
+    so Malta will be kept.
+
+    >>> s.nlargest(3).execute()
+    France    65000000
+    Italy     59000000
+    Malta       434000
+    dtype: int64
+
+    The `n` largest elements where ``n=3`` and keeping the last duplicates.
+    Brunei will be kept since it is the last with value 434000 based on
+    the index order.
+
+    >>> s.nlargest(3, keep='last').execute()
+    France      65000000
+    Italy       59000000
+    Brunei        434000
+    dtype: int64
+
+    The `n` largest elements where ``n=3`` with all duplicates kept. Note
+    that the returned Series has five elements due to the three duplicates.
+
+    >>> s.nlargest(3, keep='all').execute()
+    France      65000000
+    Italy       59000000
+    Malta         434000
+    Maldives      434000
+    Brunei        434000
+    dtype: int64
+    """
+    return _nlargest(df, n, keep=keep)