maxframe 2.0.0b2__cp37-cp37m-win32.whl → 2.3.0rc1__cp37-cp37m-win32.whl

maxframe/tensor/utils.py

@@ -167,7 +167,7 @@ def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
         except TypeError:
             pass
     # Going via an iterator directly is slower than via list comprehension.
-    axis = tuple([validate_axis(ndim, ax, argname) for ax in axis])
+    axis = tuple(validate_axis(ndim, ax, argname) for ax in axis)
     if not allow_duplicate and len(set(axis)) != len(axis):
         if argname:
             raise ValueError(f"repeated axis in `{argname}` argument")
@@ -709,8 +709,8 @@ def implement_scipy(scipy_fun_name):
     return wrapper
 
 
-def infer_scipy_dtype(scipy_fun_name):
+def infer_scipy_dtype(scipy_fun_name, **kw):
     scipy_fun = _load_scipy_func(scipy_fun_name)
     if scipy_fun is None:
         return lambda x: x
-    return infer_dtype(scipy_fun)
+    return infer_dtype(scipy_fun, **kw)

maxframe/tests/test_udf.py

@@ -0,0 +1,61 @@
+# 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 textwrap
+
+from odps import ODPS
+from odps.errors import NoSuchObject
+
+from maxframe.tests.utils import tn
+from maxframe.udf import ODPSFunction
+
+
+def test_odps_function():
+    func_body = """from odps.udf import annotate
+    @annotate("bigint->bigint")
+    class MyMul(object):
+        def evaluate(self, arg0):
+            return arg0 * 2 if arg0 is not None else None"""
+    odps_entry = ODPS.from_environments()
+    res_name = tn("test_res")
+    func_name = tn("test_odps_func")
+
+    def _cleanup():
+        try:
+            odps_entry.delete_resource(res_name + ".py")
+        except NoSuchObject:
+            pass
+        try:
+            odps_entry.delete_function(func_name)
+        except NoSuchObject:
+            pass
+
+    _cleanup()
+
+    try:
+        test_res = odps_entry.create_resource(
+            res_name + ".py", "py", fileobj=textwrap.dedent(func_body)
+        )
+        test_odps_func_obj = odps_entry.create_function(
+            func_name, class_type=f"{res_name}.MyMul", resources=[test_res]
+        )
+        func = ODPSFunction.wrap(test_odps_func_obj)
+        assert isinstance(func, ODPSFunction)
+        assert func.__name__ == func_name
+        assert func.full_function_name in (
+            f"{odps_entry.project}:{func_name}",
+            f"{odps_entry.project}:default:{func_name}",
+        )
+    finally:
+        _cleanup()

maxframe/tests/test_utils.py

@@ -31,8 +31,9 @@ import pyarrow as pa
 import pytest
 
 from .. import utils
+from ..lib.dtypes_extension import ArrowDtype
 from ..serialization import PickleContainer
-from ..utils import parse_size_to_megabytes
+from ..utils import parse_size_to_megabytes, validate_and_adjust_resource_ratio
 
 
 def test_string_conversion():
@@ -298,11 +299,11 @@ def test_estimate_pandas_size():
     s1 = pd.Series(np.random.rand(1000))
     assert utils.estimate_pandas_size(s1) == sys.getsizeof(s1)
 
-    from ..dataframe.arrays import ArrowStringArray
-
-    array = ArrowStringArray(np.random.choice(["abcd", "def", "gh"], size=(1000,)))
-    s2 = pd.Series(array)
-    assert utils.estimate_pandas_size(s2) == sys.getsizeof(s2)
+    if hasattr(pd, "ArrowDtype"):
+        arrow_array = pa.array(np.random.choice(["abcd", "def", "gh"], size=(1000,)))
+        array = pd.array(arrow_array, dtype=ArrowDtype(arrow_array.type))
+        s2 = pd.Series(array)
+        assert utils.estimate_pandas_size(s2) == sys.getsizeof(s2)
 
     s3 = pd.Series(np.random.choice(["abcd", "def", "gh"], size=(1000,)))
     assert (
@@ -366,9 +367,17 @@ def test_arrow_type_from_string():
     _assert_arrow_type_convert(pa.decimal128(10, 2))
     _assert_arrow_type_convert(pa.list_(pa.int64()))
     _assert_arrow_type_convert(pa.map_(pa.string(), pa.int64()))
+    _assert_arrow_type_convert(pa.date32())
+    _assert_arrow_type_convert(pa.date64())
     _assert_arrow_type_convert(
         pa.struct([("key", pa.string()), ("value", pa.list_(pa.int64()))])
     )
+    _assert_arrow_type_convert(
+        pa.struct([("key", pa.string(), False), ("value", pa.list_(pa.int64()))])
+    )
+    _assert_arrow_type_convert(
+        pa.struct([("key", pa.string()), ("value", pa.list_(pa.int64()), False)])
+    )
 
 
 @pytest.mark.parametrize("use_async", [False, True])
@@ -527,6 +536,7 @@ def test_numeric_inputs_with_default_units(value, default_unit):
     "input_string, expected",
     [
         # Basic binary units
+        ("1B", 1 / BYTES_PER_MIB),
         ("1KiB", BYTES_PER_KIB / BYTES_PER_MIB),
         ("5miB", 5),
         ("2giB", 2 * BYTES_PER_GIB / BYTES_PER_MIB),
@@ -571,3 +581,38 @@ def test_parse_size_to_mega_bytes_invalid_inputs(invalid_input, default_unit):
     """Test invalid inputs that should raise ValueError"""
     with pytest.raises(ValueError):  # Catch ValueError
         parse_size_to_megabytes(invalid_input, default_number_unit=default_unit)
+
+
+@pytest.mark.parametrize(
+    "udf_resources, max_memory_cpu_ratio, adjust, expected_resources, expected_adjusted, should_warn",
+    [
+        ({"other": "value"}, 4, False, {"other": "value"}, False, False),
+        ({"memory": 8}, 4, False, {"memory": 8}, False, False),
+        ({"cpu": 2}, 4, False, {"cpu": 2}, False, False),
+        ({"cpu": 2, "memory": 2}, 4, True, {"cpu": 2, "memory": 2}, False, False),
+        ({"cpu": 2, "memory": 8}, 4, False, {"cpu": 2, "memory": 8}, False, False),
+        ({"cpu": 1, "memory": 8}, 4, False, {"cpu": 1, "memory": 8}, False, False),
+        ({"cpu": 2, "memory": 8}, 4, False, {"cpu": 2, "memory": 8}, False, False),
+        ({"cpu": 2, "memory": 8}, 4, False, {"cpu": 2, "memory": 8}, False, False),
+        ({"cpu": 1, "memory": 8}, 4, True, {"cpu": 2, "memory": 8}, True, True),
+        ({"cpu": 1, "memory": 18}, 7, True, {"cpu": 3, "memory": 18}, True, True),
+        ({"cpu": 1, "memory": 7.5}, 4, True, {"cpu": 2, "memory": 7.5}, True, True),
+    ],
+)
+def test_validate_and_adjust_resource_ratio(
+    udf_resources,
+    max_memory_cpu_ratio,
+    adjust,
+    expected_resources,
+    expected_adjusted,
+    should_warn,
+    recwarn,
+):
+    result_resources, was_adjusted = validate_and_adjust_resource_ratio(
+        udf_resources, max_memory_cpu_ratio, adjust
+    )
+    assert result_resources == expected_resources
+    assert was_adjusted == expected_adjusted
+    if should_warn:
+        # check warning
+        assert len(recwarn) == 1

maxframe/tests/utils.py

@@ -16,7 +16,6 @@ import asyncio
 import contextlib
 import functools
 import hashlib
-import logging
 import os
 import queue
 import socket
@@ -233,7 +232,6 @@ def create_test_volume(vol_name, oss_config):
         yield vol_name
     finally:
         try:
-            logging.warning("Deleting test volume %s", vol_name)
             odps_entry.delete_volume(vol_name, auto_remove_dir=True, recursive=True)
         except:
             pass

maxframe/typing_.py

@@ -15,6 +15,7 @@
 from numbers import Integral
 from typing import List, TypeVar, Union
 
+import numpy as np
 import pandas as pd
 import pyarrow as pa
 
@@ -22,6 +23,7 @@ SlicesType = List[Union[None, Integral, slice]]
 
 TimeoutType = Union[int, float, None]
 
+PandasDType = Union[np.dtype, pd.api.extensions.ExtensionDtype]
 
 ArrowTableType = Union[pa.Table, pa.RecordBatch]
 PandasObjectTypes = Union[

maxframe/udf.py

@@ -13,12 +13,18 @@
 # limitations under the License.
 
 import shlex
+import sys
 from typing import Callable, List, Optional, Union
 
-from odps.models import Resource
+import numpy as np
+from odps.models import Function as ODPSFunctionObj
+from odps.models import Resource as ODPSResourceObj
 
+from .config.validators import is_positive_integer
+from .core.mode import is_mock_mode
 from .serialization import load_member
 from .serialization.serializables import (
+    AnyField,
     BoolField,
     DictField,
     FieldTypes,
@@ -27,7 +33,8 @@ from .serialization.serializables import (
     Serializable,
     StringField,
 )
-from .utils import extract_class_name, tokenize
+from .typing_ import PandasDType
+from .utils import extract_class_name, make_dtype, tokenize
 
 
 class PythonPackOptions(Serializable):
@@ -106,6 +113,7 @@ class MarkedFunction(Serializable):
     expect_resources = DictField(
         "expect_resources", FieldTypes.string, default_factory=dict
     )
+    gpu = BoolField("gpu", default=False)
 
     def __init__(self, func: Optional[Callable] = None, **kw):
         super().__init__(func=func, **kw)
@@ -120,8 +128,100 @@ class MarkedFunction(Serializable):
         return f"<MarkedFunction {self.func!r}>"
 
 
-def with_resources(*resources: Union[str, Resource], use_wrapper_class: bool = True):
-    def res_to_str(res: Union[str, Resource]) -> str:
+class ODPSFunction(Serializable):
+    __slots__ = ("_caller_type",)
+
+    full_function_name = StringField("full_function_name")
+    expect_engine = StringField("expect_engine", default=None)
+    expect_resources = DictField(
+        "expect_resources", FieldTypes.string, default_factory=dict
+    )
+    result_dtype = AnyField("result_dtype", default=None)
+
+    def __init__(
+        self,
+        func,
+        expect_engine: str = None,
+        expect_resources: dict = None,
+        dtype: PandasDType = None,
+        **kw,
+    ):
+        full_function_name = None
+        if isinstance(func, str):
+            full_function_name = func
+        elif isinstance(func, ODPSFunctionObj):
+            func_parts = [func.project.name]
+            if func.schema:
+                func_parts.append(func.schema.name)
+            func_parts.append(func.name)
+            full_function_name = ":".join(func_parts)
+        if full_function_name:
+            kw["full_function_name"] = full_function_name
+
+        if dtype is not None:
+            kw["result_dtype"] = make_dtype(dtype)
+        super().__init__(
+            expect_engine=expect_engine, expect_resources=expect_resources, **kw
+        )
+
+    @property
+    def __name__(self):
+        return self.full_function_name.rsplit(":", 1)[-1]
+
+    def _detect_caller_type(self) -> Optional[str]:
+        if hasattr(self, "_caller_type"):
+            return self._caller_type
+
+        frame = sys._getframe(1)
+        is_set = False
+        while frame.f_back:
+            f_mod = frame.f_globals.get("__name__")
+            if f_mod and f_mod.startswith("maxframe.dataframe."):
+                if f_mod.endswith(".map"):
+                    self._caller_type, is_set = "map", True
+                elif f_mod.endswith(".aggregation") or ".reduction." in f_mod:
+                    self._caller_type, is_set = "agg", True
+                if is_set:
+                    return self._caller_type
+            frame = frame.f_back
+        return None
+
+    def __call__(self, obj, *args, **kwargs):
+        caller_type = self._detect_caller_type()
+        if caller_type == "agg":
+            return self._call_aggregate(obj, *args, **kwargs)
+        raise NotImplementedError("Need to be referenced inside apply or map functions")
+
+    def _call_aggregate(self, obj, *args, **kwargs):
+        from .dataframe.core import DATAFRAME_TYPE, SERIES_TYPE
+        from .dataframe.reduction.custom_reduction import build_custom_reduction_result
+
+        if isinstance(obj, (DATAFRAME_TYPE, SERIES_TYPE)):
+            return build_custom_reduction_result(obj, self)
+        if is_mock_mode():
+            ret = obj.iloc[0]
+            if self.result_dtype:
+                if hasattr(ret, "astype"):
+                    ret = ret.astype(self.result_dtype)
+                else:  # pragma: no cover
+                    ret = np.array(ret).astype(self.result_dtype).item()
+            return ret
+        raise NotImplementedError("Need to be referenced inside apply or map functions")
+
+    def __repr__(self):
+        return f"<ODPSStoredFunction {self.full_function_name}>"
+
+    @classmethod
+    def wrap(cls, func):
+        if isinstance(func, ODPSFunctionObj):
+            return ODPSFunction(func)
+        return func
+
+
+def with_resources(
+    *resources: Union[str, ODPSResourceObj], use_wrapper_class: bool = True
+):
+    def res_to_str(res: Union[str, ODPSResourceObj]) -> str:
         if isinstance(res, str):
             return res
         res_parts = [res.project.name]
@@ -177,6 +277,8 @@ def with_running_options(
     engine: Optional[str] = None,
     cpu: Optional[int] = None,
     memory: Optional[Union[str, int]] = None,
+    gu: Optional[int] = None,
+    gu_quota: Optional[Union[str, List[str]]] = None,
     **kwargs,
 ):
     """
@@ -191,6 +293,10 @@ def with_running_options(
     memory: Optional[Union[str, int]]
         The memory to run the UDF. If it is an int, it is in GB.
         If it is a str, it is in the format of "10GiB", "30MiB", etc.
+    gu: Optional[int]
+        The GU number to run the UDF.
+    gu_quota: Optional[Union[str, List[str]]]
+        The GU quota nicknames to run the UDF. The order is the priority of the usage.
     kwargs
         Other running options.
     """
@@ -204,20 +310,37 @@ def with_running_options(
             raise TypeError("memory must be an int or str")
         if isinstance(memory, int) and memory <= 0:
             raise ValueError("memory must be greater than 0")
+    if gu is not None and gu <= 0:
+        raise ValueError("gu must be greater than 0")
+    if gu is not None and (cpu or memory):
+        raise ValueError("gu can't be specified with cpu or memory")
 
     if cpu:
         resources["cpu"] = cpu
     if memory:
         resources["memory"] = memory
 
+    if isinstance(gu_quota, str):
+        gu_quota = [gu_quota]
+
+    resources["gpu"] = gu
+    resources["gu_quota"] = gu_quota
+    use_gpu = is_positive_integer(gu)
+
     def func_wrapper(func):
-        if all(v is None for v in (engine, cpu, memory)):
+        if all(v is None for v in (engine, cpu, memory, gu, gu_quota)):
             return func
         if isinstance(func, MarkedFunction):
             func.expect_engine = engine
             func.expect_resources = resources
+            func.gpu = use_gpu
             return func
-        return MarkedFunction(func, expect_engine=engine, expect_resources=resources)
+        return MarkedFunction(
+            func,
+            expect_engine=engine,
+            expect_resources=resources,
+            gpu=use_gpu,
+        )
 
     return func_wrapper
 
@@ -225,9 +348,7 @@ def with_running_options(
 with_resource_libraries = with_resources
 
 
-def get_udf_resources(
-    func: Callable,
-) -> List[Union[Resource, str]]:
+def get_udf_resources(func: Callable) -> List[Union[ODPSResourceObj, str]]:
     return getattr(func, "resources", None) or []