maxframe 2.0.0b2__cp39-cp39-macosx_10_9_universal2.whl → 2.2.0__cp39-cp39-macosx_10_9_universal2.whl

maxframe/dataframe/sort/nsmallest.py

@@ -0,0 +1,228 @@
+# 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 _nsmallest(df, n, columns=None, keep="first"):
+    op = DataFrameSortValues(
+        output_types=[OutputType.dataframe],
+        axis=0,
+        by=columns,
+        ignore_index=False,
+        ascending=True,
+        nrows=n,
+        keep_kind=keep,
+    )
+    return op(df)
+
+
+def df_nsmallest(df, n, columns, keep="first"):
+    """
+    Return the first `n` rows ordered by `columns` in ascending order.
+
+    Return the first `n` rows with the smallest values in `columns`, in
+    ascending 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=True).head(n)``, but more
+    performant.
+
+    Parameters
+    ----------
+    n : int
+        Number of items to retrieve.
+    columns : list or str
+        Column name or names to order by.
+    keep : {'first', 'last', 'all'}, default 'first'
+        Where there are duplicate values:
+
+        - ``first`` : take the first occurrence.
+        - ``last`` : take the last occurrence.
+        - ``all`` : do not drop any duplicates, even it means
+          selecting more than `n` items.
+
+    Returns
+    -------
+    DataFrame
+
+    See Also
+    --------
+    DataFrame.nlargest : Return the first `n` rows ordered by `columns` in
+        descending order.
+    DataFrame.sort_values : Sort DataFrame by the values.
+    DataFrame.head : Return the first `n` rows without re-ordering.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({'population': [59000000, 65000000, 434000,
+    ...                                   434000, 434000, 337000, 337000,
+    ...                                   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         337000      182      NR
+    Tuvalu         11300       38      TV
+    Anguilla       11300      311      AI
+
+    In the following example, we will use ``nsmallest`` to select the
+    three rows having the smallest values in column "population".
+
+    >>> df.nsmallest(3, 'population').execute()
+              population    GDP alpha-2
+    Tuvalu         11300     38      TV
+    Anguilla       11300    311      AI
+    Iceland       337000  17036      IS
+
+    When using ``keep='last'``, ties are resolved in reverse order:
+
+    >>> df.nsmallest(3, 'population', keep='last').execute()
+              population  GDP alpha-2
+    Anguilla       11300  311      AI
+    Tuvalu         11300   38      TV
+    Nauru         337000  182      NR
+
+    When using ``keep='all'``, all duplicate items are maintained:
+
+    >>> df.nsmallest(3, 'population', keep='all').execute()
+              population    GDP alpha-2
+    Tuvalu         11300     38      TV
+    Anguilla       11300    311      AI
+    Iceland       337000  17036      IS
+    Nauru         337000    182      NR
+
+    To order by the smallest values in column "population" and then "GDP", we can
+    specify multiple columns like in the next example.
+
+    >>> df.nsmallest(3, ['population', 'GDP']).execute()
+              population  GDP alpha-2
+    Tuvalu         11300   38      TV
+    Anguilla       11300  311      AI
+    Nauru         337000  182      NR
+    """
+    return _nsmallest(df, n, columns, keep=keep)
+
+
+def series_nsmallest(df, n, keep="first"):
+    """
+    Return the smallest `n` elements.
+
+    Parameters
+    ----------
+    n : int, default 5
+        Return this many ascending 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` smallest values in the Series, sorted in increasing order.
+
+    See Also
+    --------
+    Series.nlargest: Get the `n` largest elements.
+    Series.sort_values: Sort Series by values.
+    Series.head: Return the first `n` rows.
+
+    Notes
+    -----
+    Faster than ``.sort_values().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,
+    ...                         "Brunei": 434000, "Malta": 434000,
+    ...                         "Maldives": 434000, "Iceland": 337000,
+    ...                         "Nauru": 11300, "Tuvalu": 11300,
+    ...                         "Anguilla": 11300, "Montserrat": 5200}
+    >>> s = md.Series(countries_population)
+    >>> s.execute()
+    Italy       59000000
+    France      65000000
+    Brunei        434000
+    Malta         434000
+    Maldives      434000
+    Iceland       337000
+    Nauru          11300
+    Tuvalu         11300
+    Anguilla       11300
+    Montserrat      5200
+    dtype: int64
+
+    The `n` smallest elements where ``n=5`` by default.
+
+    >>> s.nsmallest().execute()
+    Montserrat    5200
+    Nauru        11300
+    Tuvalu       11300
+    Anguilla     11300
+    Iceland     337000
+    dtype: int64
+
+    The `n` smallest elements where ``n=3``. Default `keep` value is
+    'first' so Nauru and Tuvalu will be kept.
+
+    >>> s.nsmallest(3).execute()
+    Montserrat   5200
+    Nauru       11300
+    Tuvalu      11300
+    dtype: int64
+
+    The `n` smallest elements where ``n=3`` and keeping the last
+    duplicates. Anguilla and Tuvalu will be kept since they are the last
+    with value 11300 based on the index order.
+
+    >>> s.nsmallest(3, keep='last').execute()
+    Montserrat   5200
+    Anguilla    11300
+    Tuvalu      11300
+    dtype: int64
+
+    The `n` smallest elements where ``n=3`` with all duplicates kept. Note
+    that the returned Series has four elements due to the three duplicates.
+
+    >>> s.nsmallest(3, keep='all').execute()
+    Montserrat   5200
+    Nauru       11300
+    Tuvalu      11300
+    Anguilla    11300
+    dtype: int64
+    """
+    return _nsmallest(df, n, keep=keep)

maxframe/dataframe/statistics/__init__.py

@@ -16,15 +16,15 @@
 def _install():
     from ..core import DATAFRAME_TYPE, SERIES_TYPE
     from .corr import df_corr, df_corrwith, series_autocorr, series_corr
-    from .quantile import quantile_dataframe, quantile_series
+    from .quantile import dataframe_quantile, series_quantile
 
     for t in SERIES_TYPE:
-        t.quantile = quantile_series
+        t.quantile = series_quantile
         t.corr = series_corr
         t.autocorr = series_autocorr
 
     for t in DATAFRAME_TYPE:
-        t.quantile = quantile_dataframe
+        t.quantile = dataframe_quantile
         t.corr = df_corr
         t.corrwith = df_corrwith
 

maxframe/dataframe/statistics/corr.py

@@ -34,6 +34,7 @@ class DataFrameCorr(DataFrameOperator, DataFrameOperatorMixin):
     min_periods = Int32Field("min_periods", default=None)
     axis = Int32Field("axis", default=None)
     drop = BoolField("drop", default=None)
+    ddof = Int32Field("ddof", default=0)
 
     @classmethod
     def _set_inputs(cls, op: "DataFrameCorr", inputs: List[EntityData]):

maxframe/dataframe/statistics/quantile.py

@@ -213,7 +213,7 @@ class DataFrameQuantile(DataFrameOperator, DataFrameOperatorMixin):
             return self._call_series(a, inputs)
 
 
-def quantile_series(series, q=0.5, interpolation="linear"):
+def series_quantile(series, q=0.5, interpolation="linear"):
     """
     Return value at the given quantile.
 
@@ -268,7 +268,7 @@ def quantile_series(series, q=0.5, interpolation="linear"):
     return op(series, q_input=q_input)
 
 
-def quantile_dataframe(df, q=0.5, axis=0, numeric_only=True, interpolation="linear"):
+def dataframe_quantile(df, q=0.5, axis=0, numeric_only=True, interpolation="linear"):
     # FIXME: Timedelta not support. Data invalid: ODPS-0010000:InvalidArgument:duration[ns] is not equal to string
     """
     Return values at the given quantile over requested axis.

maxframe/dataframe/tests/test_typing.py

@@ -0,0 +1,104 @@
+# 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.
+
+import numpy as np
+import pandas as pd
+
+from ...core import OutputType
+from ..typing_ import get_function_output_meta
+
+
+def test_dataframe_type_annotation():
+    def func() -> pd.DataFrame[int]:
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta is not None
+    assert meta.output_type.name == "dataframe"
+    assert len(meta.dtypes) == 1
+    assert meta.dtypes[0] == np.dtype(int)
+
+    def func() -> pd.DataFrame[{"col1": int, "col2": float}]:  # noqa: F821
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta is not None
+    assert meta.output_type.name == "dataframe"
+    assert len(meta.dtypes) == 2
+    assert meta.dtypes[0] == np.dtype(int)
+    assert meta.dtypes[1] == np.dtype(float)
+
+    def func() -> pd.DataFrame[str, {"col1": int, "col2": float}]:  # noqa: F821
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta is not None
+    assert meta.output_type.name == "dataframe"
+    assert len(meta.dtypes) == 2
+    assert meta.index_value.value.dtype == np.dtype("O")
+    assert list(meta.dtypes.index) == ["col1", "col2"]
+    assert list(meta.dtypes) == [np.dtype(int), np.dtype(float)]
+
+
+def test_series_type_annotation():
+    def func() -> pd.Series[np.str_]:
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta is not None
+    assert meta.output_type == OutputType.series
+    assert meta.dtype == np.dtype(np.str_)
+
+    def func() -> pd.Series[("idx_name", str), ("series_name", np.int64)]:  # noqa: F821
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta is not None
+    assert meta.output_type == OutputType.series
+    assert meta.name == "series_name"
+    assert meta.dtype == np.dtype(np.int64)
+    assert meta.index_value.value._name == "idx_name"
+    assert meta.index_value.value.dtype == np.dtype("O")
+
+
+def test_index_type_annotation():
+    def func() -> pd.Index[np.int64]:
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta is not None
+    assert meta.output_type == OutputType.index
+
+    def func() -> pd.Index[[("ix1", str), ("ix2", np.int64)]]:  # noqa: F821
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta is not None
+    assert meta.output_type == OutputType.index
+    assert meta.index_value.value.names == ["ix1", "ix2"]
+    assert list(meta.index_value.value.dtypes) == [np.dtype("O"), np.dtype("int64")]
+
+
+def test_function_output_meta_corner_cases():
+    def func():
+        pass
+
+    assert get_function_output_meta(func) is None
+    assert get_function_output_meta("non-func-obj") is None
+
+    def func() -> int:
+        pass
+
+    meta = get_function_output_meta(func)
+    assert meta.dtype == np.dtype("int64")

maxframe/dataframe/tests/test_utils.py

@@ -17,9 +17,16 @@ import pandas as pd
 import pyarrow as pa
 import pytest
 
-from ...udf import MarkedFunction, with_python_requirements, with_resources
+from ...config import option_context
+from ...core.operator import Operator
+from ...udf import (
+    MarkedFunction,
+    with_python_requirements,
+    with_resources,
+    with_running_options,
+)
 from ...utils import ARROW_DTYPE_NOT_SUPPORTED
-from ..utils import _generate_value, pack_func_args
+from ..utils import _generate_value, copy_func_scheduling_hints, pack_func_args
 
 try:
     from pandas import ArrowDtype
@@ -84,6 +91,20 @@ def test_pack_function(df1):
             [(np.int32(1), "1")],
         ),
         (pa.map_(pa.int32(), pa.string()), 1, [(np.int32(1), "1")]),
+        (
+            ArrowDtype(
+                pa.struct([pa.field("a", pa.int32()), pa.field("b", pa.string())])
+            )
+            if ArrowDtype
+            else None,
+            1,
+            {"a": np.int32(1), "b": "1"},
+        ),
+        (
+            pa.struct([pa.field("a", pa.int32()), pa.field("b", pa.string())]),
+            1,
+            {"a": np.int32(1), "b": "1"},
+        ),
         (pa.int32(), 1, np.int32(1)),
         (np.datetime64, "2023-01-01", pd.Timestamp("2023-01-01")),
         (np.timedelta64, "1D", pd.Timedelta("1D")),
@@ -99,3 +120,46 @@ def test_pack_function(df1):
 def test_generate_value(dtype, fill_value, expected):
     result = _generate_value(dtype, fill_value)
     assert result == expected
+
+
+def test_copy_func_scheduling_hints():
+    # Test with a regular function (no scheduling hints)
+    with option_context() as options:
+        options.function.default_running_options = {}  # No default options
+
+        def regular_func(x):
+            return x + 1
+
+        op1 = Operator()
+        copy_func_scheduling_hints(regular_func, op1)
+        # Should not set any attributes since regular function has no hints
+        assert not hasattr(op1, "expect_engine") or op1.expect_engine is None
+        assert not hasattr(op1, "expect_resources") or op1.expect_resources is None
+        assert not hasattr(op1, "gpu") or op1.gpu is None
+
+    # Test with MarkedFunction with scheduling hints
+
+    @with_running_options(engine="DPE", cpu=4, memory="8GiB")
+    def marked_func(x):
+        return x + 1
+
+    op2 = Operator()
+    copy_func_scheduling_hints(marked_func, op2)
+    assert op2.expect_engine == "DPE"
+    # The expect_resources will include default values for gpu and gu_quota
+    expected_resources = {"cpu": 4, "memory": "8GiB", "gpu": 0, "gu_quota": None}
+    assert op2.expect_resources == expected_resources
+
+    # Test with MarkedFunction with GPU
+    @with_running_options(gu=2)
+    def gpu_func(x):
+        return x + 1
+
+    op3 = Operator()
+    copy_func_scheduling_hints(gpu_func, op3)
+    assert op3.gpu is True
+    # The expect_resources will include the gu value and default values
+    # System has default options: {'cpu': 1, 'memory': '4GiB', 'gpu': 0}
+    # The with_running_options decorator will override the gpu value with the gu value
+    expected_resources = {"gpu": 2, "gu_quota": None, "cpu": 1, "memory": "4GiB"}
+    assert op3.expect_resources == expected_resources

maxframe/dataframe/typing_.py

@@ -0,0 +1,185 @@
+# 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.
+
+import dataclasses
+import functools
+import inspect
+from typing import Any, Callable, Generic, List, Optional, TypeVar
+
+import pandas as pd
+
+from ..core import OutputType
+from ..typing_ import PandasDType
+from ..utils import make_dtype
+from .utils import InferredDataFrameMeta, parse_index
+
+# TypeVars
+T = TypeVar("T")
+
+
+@dataclasses.dataclass
+class _FieldDef:
+    name: Any
+    dtype: PandasDType
+
+
+def _item_to_field_def(item_):
+    if isinstance(item_, tuple):
+        tp = make_dtype(item_[1])
+        return _FieldDef(name=item_[0], dtype=tp)
+    else:
+        tp = make_dtype(item_)
+        return _FieldDef(name=None, dtype=tp)
+
+
+class IndexType:
+    def __init__(self, index_fields: List[_FieldDef]):
+        self.index_fields = index_fields
+
+    def __repr__(self):
+        return f"IndexType({[f.dtype for f in self.index_fields]})"
+
+    @classmethod
+    def from_getitem_args(cls, item) -> "IndexType":
+        if isinstance(item, (dict, pd.Series)):
+            item = list(item.items())
+
+        if isinstance(item, list):
+            return IndexType([_item_to_field_def(tp) for tp in item])
+        else:
+            return IndexType([_item_to_field_def(item)])
+
+
+class SeriesType(Generic[T]):
+    def __init__(
+        self, index_fields: Optional[List[_FieldDef]], name_and_dtype: _FieldDef
+    ):
+        self.index_fields = index_fields
+        self.name_and_dtype = name_and_dtype
+
+    def __repr__(self) -> str:
+        return "SeriesType[{}]".format(self.name_and_dtype.dtype)
+
+    @classmethod
+    def from_getitem_args(cls, item) -> "SeriesType":
+        if not isinstance(item, tuple):
+            item = (item,)
+        if len(item) == 1:
+            tp = _item_to_field_def(item[0])
+            return SeriesType(None, tp)
+        else:
+            tp = _item_to_field_def(item[1])
+            idx_fields = IndexType.from_getitem_args(item[0]).index_fields
+            return SeriesType(idx_fields, tp)
+
+
+class DataFrameType:
+    def __init__(
+        self,
+        index_fields: Optional[List[_FieldDef]],
+        data_fields: List[_FieldDef],
+    ):
+        self.index_fields = index_fields
+        self.data_fields = data_fields
+
+    def __repr__(self) -> str:
+        types = [field.dtype for field in self.data_fields]
+        return f"DataFrameType[{types}]"
+
+    @classmethod
+    def from_getitem_args(cls, item) -> "DataFrameType":
+        if not isinstance(item, tuple):
+            item = (item,)
+        fields = IndexType.from_getitem_args(item[-1]).index_fields
+        if len(item) == 1:
+            return DataFrameType(None, fields)
+        else:
+            idx_fields = IndexType.from_getitem_args(item[0]).index_fields
+            return DataFrameType(idx_fields, fields)
+
+
+def get_function_output_meta(
+    func: Callable, df_obj=None
+) -> Optional[InferredDataFrameMeta]:
+    try:
+        func_argspec = inspect.getfullargspec(func)
+        ret_type = (func_argspec.annotations or {}).get("return")
+        if ret_type is None:
+            return None
+    except:
+        return None
+
+    dtypes = dtype = name = None
+    index_fields = None
+    if isinstance(ret_type, DataFrameType):
+        output_type = OutputType.dataframe
+        dtypes = pd.Series(
+            [fd.dtype for fd in ret_type.data_fields],
+            index=[fd.name for fd in ret_type.data_fields],
+        )
+        index_fields = ret_type.index_fields
+    elif isinstance(ret_type, SeriesType):
+        output_type = OutputType.series
+        dtype = ret_type.name_and_dtype.dtype
+        name = ret_type.name_and_dtype.name
+        index_fields = ret_type.index_fields
+    elif isinstance(ret_type, IndexType):
+        output_type = OutputType.index
+        index_fields = ret_type.index_fields
+    else:
+        output_type = OutputType.scalar
+        try:
+            dtype = make_dtype(ret_type)
+        except:
+            return None
+
+    if index_fields is not None:
+        if len(index_fields) == 1:
+            mock_idx = pd.Index(
+                [], dtype=index_fields[0].dtype, name=index_fields[0].name
+            )
+        else:
+            col_names = [index_field.name for index_field in index_fields]
+            col_dtypes = pd.Series(
+                [index_field.dtype for index_field in index_fields], index=col_names
+            )
+            mock_df = pd.DataFrame([], columns=col_names).astype(col_dtypes)
+            mock_idx = pd.MultiIndex.from_frame(mock_df)
+        index_value = parse_index(mock_idx, df_obj, store_data=False)
+    else:
+        index_value = None
+
+    return InferredDataFrameMeta(
+        output_type=output_type,
+        index_value=index_value,
+        dtypes=dtypes,
+        dtype=dtype,
+        name=name,
+    )
+
+
+def register_pandas_typing_funcs():
+    def _cls_getitem_func(cls, item, type_cls):
+        return type_cls.from_getitem_args(item)
+
+    for pd_cls, type_cls in [
+        (pd.DataFrame, DataFrameType),
+        (pd.Series, SeriesType),
+        (pd.Index, IndexType),
+    ]:
+        if hasattr(pd_cls, "__class_getitem__"):  # pragma: no cover
+            continue
+        pd_cls.__class_getitem__ = classmethod(
+            functools.partial(_cls_getitem_func, type_cls=type_cls)
+        )