maxframe 1.3.0__cp311-cp311-macosx_10_9_universal2.whl → 2.0.0__cp311-cp311-macosx_10_9_universal2.whl

maxframe/tensor/statistics/average.py

@@ -0,0 +1,143 @@
+# 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
+
+from ..datasource import tensor as astensor
+from ..misc.broadcast_to import broadcast_to
+from ..misc.swapaxes import swapaxes
+
+
+def average(a, axis=None, weights=None, returned=False):
+    """
+    Compute the weighted average along the specified axis.
+
+    Parameters
+    ----------
+    a : array_like
+        Tensor containing data to be averaged. If `a` is not a tensor, a
+        conversion is attempted.
+    axis : None or int or tuple of ints, optional
+        Axis or axes along which to average `a`.  The default,
+        axis=None, will average over all of the elements of the input tensor.
+        If axis is negative it counts from the last to the first axis.
+
+        If axis is a tuple of ints, averaging is performed on all of the axes
+        specified in the tuple instead of a single axis or all the axes as
+        before.
+    weights : array_like, optional
+        A tensor of weights associated with the values in `a`. Each value in
+        `a` contributes to the average according to its associated weight.
+        The weights tensor can either be 1-D (in which case its length must be
+        the size of `a` along the given axis) or of the same shape as `a`.
+        If `weights=None`, then all data in `a` are assumed to have a
+        weight equal to one.
+    returned : bool, optional
+        Default is `False`. If `True`, the tuple (`average`, `sum_of_weights`)
+        is returned, otherwise only the average is returned.
+        If `weights=None`, `sum_of_weights` is equivalent to the number of
+        elements over which the average is taken.
+
+
+    Returns
+    -------
+    average, [sum_of_weights] : tensor_type or double
+        Return the average along the specified axis. When returned is `True`,
+        return a tuple with the average as the first element and the sum
+        of the weights as the second element. The return type is `Float`
+        if `a` is of integer type, otherwise it is of the same type as `a`.
+        `sum_of_weights` is of the same type as `average`.
+
+    Raises
+    ------
+    ZeroDivisionError
+        When all weights along axis are zero. See `numpy.ma.average` for a
+        version robust to this type of error.
+    TypeError
+        When the length of 1D `weights` is not the same as the shape of `a`
+        along axis.
+
+    See Also
+    --------
+    mean
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> data = list(range(1,5))
+    >>> data
+    [1, 2, 3, 4]
+    >>> mt.average(data).execute()
+    2.5
+    >>> mt.average(range(1,11), weights=range(10,0,-1)).execute()
+    4.0
+
+    >>> data = mt.arange(6).reshape((3,2))
+    >>> data.execute()
+    array([[0, 1],
+           [2, 3],
+           [4, 5]])
+    >>> mt.average(data, axis=1, weights=[1./4, 3./4]).execute()
+    array([ 0.75,  2.75,  4.75])
+    >>> mt.average(data, weights=[1./4, 3./4]).execute()
+    Traceback (most recent call last):
+    ...
+    TypeError: Axis must be specified when shapes of a and weights differ.
+
+    """
+    from ..arithmetic import multiply, truediv
+
+    a = astensor(a)
+
+    if weights is None:
+        avg = a.mean(axis)
+        scl = avg.dtype.type(a.size / avg.size)
+    else:
+        wgt = astensor(weights)
+
+        if issubclass(a.dtype.type, (np.integer, np.bool_)):
+            result_dtype = np.result_type(a.dtype, wgt.dtype, "f8")
+        else:
+            result_dtype = np.result_type(a.dtype, wgt.dtype)
+
+        # sanity checks
+        if a.shape != wgt.shape:
+            if axis is None:
+                raise TypeError(
+                    "Axis must be specified when shapes of a and weights differ."
+                )
+            if wgt.ndim != 1:
+                raise TypeError(
+                    "1D weights expected when shapes of a and weights differ."
+                )
+            if wgt.shape[0] != a.shape[axis]:
+                raise ValueError(
+                    "Length of weights not compatible with specified axis."
+                )
+
+            # setup wgt to broadcast along axis
+            wgt = broadcast_to(wgt, (a.ndim - 1) * (1,) + wgt.shape)
+            wgt = swapaxes(wgt, -1, axis)
+
+        scl = wgt.sum(axis=axis, dtype=result_dtype)
+        with np.errstate(divide="raise"):
+            avg = truediv(multiply(a, wgt, dtype=result_dtype).sum(axis), scl)
+
+    if returned:
+        if scl.shape != avg.shape:
+            scl = broadcast_to(scl, avg.shape)
+        return avg, scl
+    else:
+        return avg

maxframe/tensor/statistics/bincount.py

@@ -0,0 +1,133 @@
+# 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 typing import List, Optional
+
+import numpy as np
+
+from ... import opcodes
+from ...serialization.serializables import Int64Field, ReferenceField
+from ...typing_ import EntityType
+from ..datasource import tensor as astensor
+from ..operators import TensorMapReduceOperator, TensorOperatorMixin
+
+_DEFAULT_CHUNK_SIZE_LIMIT = 1e8
+
+
+class TensorBinCount(TensorMapReduceOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.BINCOUNT
+
+    weights = ReferenceField("weights", default=None)
+    minlength: Optional[int] = Int64Field("minlength", default=0)
+    chunk_size_limit: int = Int64Field("chunk_size_limit")
+
+    chunk_count: Optional[int] = Int64Field("chunk_count")
+    tileable_right_bound: Optional[int] = Int64Field("tileable_right_bound")
+
+    def __call__(self, x, weights=None):
+        inputs = [x]
+        self.weights = weights
+        dtype = np.dtype(np.int_)
+        if weights is not None:
+            inputs.append(weights)
+            dtype = weights.dtype
+        return self.new_tensor(inputs, dtype=dtype, shape=(np.nan,))
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorBinCount", inputs: List[EntityType]):
+        super()._set_inputs(op, inputs)
+        if len(inputs) > 1:
+            op.weights = inputs[1]
+
+
+def bincount(x, weights=None, minlength=0, chunk_size_limit=None):
+    """
+    Count number of occurrences of each value in array of non-negative ints.
+
+    The number of bins (of size 1) is one larger than the largest value in
+    `x`. If `minlength` is specified, there will be at least this number
+    of bins in the output array (though it will be longer if necessary,
+    depending on the contents of `x`).
+    Each bin gives the number of occurrences of its index value in `x`.
+    If `weights` is specified the input array is weighted by it, i.e. if a
+    value ``n`` is found at position ``i``, ``out[n] += weight[i]`` instead
+    of ``out[n] += 1``.
+
+    Parameters
+    ----------
+    x : tensor or array_like, 1 dimension, nonnegative ints
+        Input array.
+    weights : tensor or array_like, optional
+        Weights, array of the same shape as `x`.
+    minlength : int, optional
+        A minimum number of bins for the output array.
+
+    Returns
+    -------
+    out : tensor of ints
+        The result of binning the input array.
+        The length of `out` is equal to ``np.amax(x)+1``.
+
+    Raises
+    ------
+    ValueError
+        If the input is not 1-dimensional, or contains elements with negative
+        values, or if `minlength` is negative.
+    TypeError
+        If the type of the input is float or complex.
+
+    See Also
+    --------
+    histogram, digitize, unique
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> mt.bincount(mt.arange(5)).execute()
+    array([1, 1, 1, 1, 1])
+    >>> mt.bincount(mt.tensor([0, 1, 1, 3, 2, 1, 7])).execute()
+    array([1, 3, 1, 1, 0, 0, 0, 1])
+
+    The input array needs to be of integer dtype, otherwise a
+    TypeError is raised:
+
+    >>> mt.bincount(mt.arange(5, dtype=float)).execute()
+    Traceback (most recent call last):
+      ....execute()
+    TypeError: Cannot cast array data from dtype('float64') to dtype('int64')
+    according to the rule 'safe'
+
+    A possible use of ``bincount`` is to perform sums over
+    variable-size chunks of an array, using the ``weights`` keyword.
+
+    >>> w = mt.array([0.3, 0.5, 0.2, 0.7, 1., -0.6]) # weights
+    >>> x = mt.array([0, 1, 1, 2, 2, 2])
+    >>> mt.bincount(x, weights=w).execute()
+    array([ 0.3,  0.7,  1.1])
+    """
+    x = astensor(x)
+    weights = astensor(weights) if weights is not None else None
+
+    if not np.issubdtype(x.dtype, np.int_):
+        raise TypeError(f"Cannot cast array data from {x.dtype} to {np.dtype(np.int_)}")
+    if x.ndim != 1:
+        raise ValueError("'x' must be 1 dimension")
+    if minlength < 0:
+        raise ValueError("'minlength' must not be negative")
+
+    chunk_size_limit = (
+        chunk_size_limit if chunk_size_limit is not None else _DEFAULT_CHUNK_SIZE_LIMIT
+    )
+    op = TensorBinCount(minlength=minlength, chunk_size_limit=chunk_size_limit)
+    return op(x, weights=weights)

maxframe/tensor/statistics/quantile.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 
 from collections.abc import Iterable
+from typing import List
 
 import numpy as np
 
 from ... import opcodes
-from ...core import ENTITY_TYPE
+from ...core import ENTITY_TYPE, EntityData
 from ...serialization.serializables import AnyField, BoolField, KeyField, StringField
 from ..core import TENSOR_TYPE, TensorOrder
 from ..datasource import tensor as astensor
@@ -57,13 +58,14 @@ class TensorQuantile(TensorOperator, TensorOperatorMixin):
         self.q_error_msg = kw.pop("q_error_msg", q_error_msg)
         super().__init__(**kw)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.a = self._inputs[0]
-        if isinstance(self.q, TENSOR_TYPE):
-            self.q = self._inputs[1]
-        if isinstance(self.out, TENSOR_TYPE):
-            self.out = self._inputs[-1]
+    @classmethod
+    def _set_inputs(cls, op: "TensorQuantile", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.a = op._inputs[0]
+        if isinstance(op.q, TENSOR_TYPE):
+            op.q = op._inputs[1]
+        if isinstance(op.out, TENSOR_TYPE):
+            op.out = op._inputs[-1]
 
     def __call__(self, a, q=None, out=None):
         shape = [self.q.size] if self.q.ndim > 0 else []

maxframe/tensor/ufunc/__init__.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

maxframe/tensor/ufunc/ufunc.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

maxframe/tensor/utils.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import inspect
 import itertools
 import operator
@@ -669,10 +668,22 @@ def fetch_corner_data(tensor, session=None):
         return tensor.fetch(session=session)
 
 
-def implement_scipy(scipy_fun):
+def _load_scipy_func(func_name: str):
+    try:
+        assert func_name.startswith("scipy.")
+        mod_name, func_name = func_name.rsplit(".", 1)
+        mod = importlib.import_module(mod_name)
+        return getattr(mod, func_name)
+    except (AttributeError, ImportError):
+        return None
+
+
+def implement_scipy(scipy_fun_name):
     import re
     import textwrap
 
+    scipy_fun = _load_scipy_func(scipy_fun_name)
+
     def wrapper(fun):
         if scipy_fun is None:
             return None
@@ -696,3 +707,10 @@ def implement_scipy(scipy_fun):
         return fun
 
     return wrapper
+
+
+def infer_scipy_dtype(scipy_fun_name):
+    scipy_fun = _load_scipy_func(scipy_fun_name)
+    if scipy_fun is None:
+        return lambda x: x
+    return infer_dtype(scipy_fun)

maxframe/tests/test_protocol.py

@@ -20,8 +20,8 @@ import pytest
 from ..lib import wrapped_pickle
 from ..protocol import (
     DagInfo,
-    DagStatus,
     ErrorInfo,
+    ExecutionStatus,
     ODPSTableResultInfo,
     ODPSVolumeResultInfo,
     ResultInfo,
@@ -127,7 +127,7 @@ def test_dag_info_json_serialize():
     info = DagInfo(
         session_id="test_session_id",
         dag_id="test_dag_id",
-        status=DagStatus.FAILED,
+        status=ExecutionStatus.FAILED,
         progress=0.65,
         tileable_to_result_infos={
             "tileable_key": ODPSTableResultInfo(full_table_name="table_name")
@@ -153,7 +153,7 @@ def test_session_info_json_serialize():
     dag_info = DagInfo(
         session_id="test_session_id",
         dag_id="test_dag_id",
-        status=DagStatus.RUNNING,
+        status=ExecutionStatus.RUNNING,
         progress=0.65,
     )
     info = SessionInfo(

maxframe/tests/test_utils.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,13 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import copy
+import functools
 import os
 import shutil
 import sys
 import tempfile
 import textwrap
 import time
+from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
 from functools import partial
 
@@ -29,6 +31,8 @@ import pyarrow as pa
 import pytest
 
 from .. import utils
+from ..serialization import PickleContainer
+from ..utils import parse_size_to_megabytes
 
 
 def test_string_conversion():
@@ -170,6 +174,9 @@ def test_tokenize():
     assert utils.tokenize(partial_f) == utils.tokenize(copy.deepcopy(partial_f))
     assert utils.tokenize(partial_f) != utils.tokenize(partial_f2)
 
+    container = PickleContainer([b"abcd", b"efgh"])
+    assert utils.tokenize(container) == utils.tokenize(copy.deepcopy(container))
+
 
 def test_lazy_import():
     old_sys_path = sys.path
@@ -362,3 +369,205 @@ def test_arrow_type_from_string():
     _assert_arrow_type_convert(
         pa.struct([("key", pa.string()), ("value", pa.list_(pa.int64()))])
     )
+
+
+@pytest.mark.parametrize("use_async", [False, True])
+async def test_call_with_retry(use_async):
+    retry_idx_list = [0]
+
+    def sync_func(delay=0):
+        if delay:
+            time.sleep(delay)
+        if retry_idx_list[0] < 3:
+            retry_idx_list[0] += 1
+            raise ValueError
+
+    async def async_func(delay=0):
+        if delay:
+            await asyncio.sleep(delay)
+        if retry_idx_list[0] < 3:
+            retry_idx_list[0] += 1
+            raise ValueError
+
+    func = async_func if use_async else sync_func
+
+    async def wait_coro(res):
+        if asyncio.iscoroutine(res):
+            return await res
+        return res
+
+    # test cases for retry times
+    with pytest.raises(ValueError):
+        retry_idx_list[0] = 0
+        await wait_coro(
+            utils.call_with_retry(func, retry_times=1, exc_type=(TypeError, ValueError))
+        )
+    assert retry_idx_list[0] == 2
+
+    retry_idx_list[0] = 0
+    await wait_coro(
+        utils.call_with_retry(func, retry_times=3, exc_type=(TypeError, ValueError))
+    )
+    assert retry_idx_list[0] == 3
+
+    retry_idx_list[0] = 0
+    exc_info = await wait_coro(
+        utils.call_with_retry(
+            func, retry_times=1, exc_type=(TypeError, ValueError), no_raise=True
+        )
+    )
+    assert isinstance(exc_info[1], ValueError)
+    assert retry_idx_list[0] == 2
+
+    delay_func = functools.partial(func, delay=0.5)
+    with pytest.raises(ValueError):
+        retry_idx_list[0] = 0
+        await wait_coro(
+            utils.call_with_retry(delay_func, retry_times=None, retry_timeout=0.7)
+        )
+    assert retry_idx_list[0] == 2
+
+    retry_idx_list[0] = 0
+    await wait_coro(
+        utils.call_with_retry(delay_func, retry_times=None, retry_timeout=2.2)
+    )
+    assert retry_idx_list[0] == 3
+
+    retry_idx_list[0] = 0
+    exc_info = await wait_coro(
+        utils.call_with_retry(
+            delay_func, retry_times=None, retry_timeout=0.7, no_raise=True
+        )
+    )
+    assert isinstance(exc_info[1], ValueError)
+    assert retry_idx_list[0] == 2
+
+
+def test_debug_to_thread():
+    class MixinTestCls(utils.ToThreadMixin):
+        async def run_in_coro(self):
+            await self.to_thread(time.sleep, 0.5)
+            await self.to_thread(functools.partial(time.sleep), 0.5)
+
+    def thread_body():
+        loop = asyncio.new_event_loop()
+        loop.set_debug(True)
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(MixinTestCls().run_in_coro())
+
+    tpe = ThreadPoolExecutor(max_workers=1)
+    tpe.submit(thread_body).result()
+    tpe.shutdown()
+
+
+@pytest.mark.parametrize(
+    "val, expected, make_series",
+    [
+        (int, np.dtype(int), True),
+        ([int], [np.dtype(int)], False),
+        ([int], pd.Series([np.dtype(int)]), True),
+        (np.dtype("float64"), np.dtype("float64"), True),
+        ("category", "category", True),
+        ("string", "string", True),
+        (pd.Timestamp, np.dtype("datetime64[ns]"), True),
+        (pd.Timedelta, np.dtype("timedelta64[ns]"), True),
+        ({"col": float}, {"col": np.dtype(float)}, False),
+        ({"col": float}, pd.Series({"col": np.dtype(float)}), True),
+        (
+            pd.Series([float], index=["col"]),
+            pd.Series([np.dtype(float)], index=["col"]),
+            True,
+        ),
+    ],
+)
+def test_make_dtypes(val, expected, make_series):
+    result = utils.make_dtypes(val, make_series=make_series)
+    if isinstance(expected, pd.Series):
+        pd.testing.assert_series_equal(result, expected)
+    else:
+        assert result == expected
+
+
+# Define conversion constants
+BYTES_PER_KIB = 1024
+BYTES_PER_MIB = 1024**2
+BYTES_PER_GIB = 1024**3
+BYTES_PER_TIB = 1024**4
+
+BYTES_PER_KB = 1000
+BYTES_PER_MB = 1000**2
+BYTES_PER_GB = 1000**3
+BYTES_PER_TB = 1000**4
+
+
+# Test numeric values with default units
+@pytest.mark.parametrize("value", [0, 1, 4, 5, 0.01, 0.1, 0.5, 1.5])
+@pytest.mark.parametrize(
+    "default_unit", ["KiB", "MiB", "GiB", "TiB", "KB", "MB", "GB", "TB"]
+)
+def test_numeric_inputs_with_default_units(value, default_unit):
+    """Test numeric inputs with various default units"""
+    # Define unit conversions to MiB for reusability
+    unit_to_mib_factor = {
+        "KiB": BYTES_PER_KIB / BYTES_PER_MIB,
+        "MiB": 1,
+        "GiB": BYTES_PER_GIB / BYTES_PER_MIB,
+        "TiB": BYTES_PER_TIB / BYTES_PER_MIB,
+        "KB": BYTES_PER_KB / BYTES_PER_MIB,
+        "MB": BYTES_PER_MB / BYTES_PER_MIB,
+        "GB": BYTES_PER_GB / BYTES_PER_MIB,
+        "TB": BYTES_PER_TB / BYTES_PER_MIB,
+    }
+    expected = value * unit_to_mib_factor[default_unit]
+    result = parse_size_to_megabytes(value, default_number_unit=default_unit)
+    assert pytest.approx(result) == expected
+
+
+@pytest.mark.parametrize(
+    "input_string, expected",
+    [
+        # Basic binary units
+        ("1KiB", BYTES_PER_KIB / BYTES_PER_MIB),
+        ("5miB", 5),
+        ("2giB", 2 * BYTES_PER_GIB / BYTES_PER_MIB),
+        ("0.1TiB", 0.1 * BYTES_PER_TIB / BYTES_PER_MIB),
+        # Basic decimal units
+        ("1KB", BYTES_PER_KB / BYTES_PER_MIB),
+        ("10MB", 10 * BYTES_PER_MB / BYTES_PER_MIB),
+        ("0.5GB", 0.5 * BYTES_PER_GB / BYTES_PER_MIB),
+        ("0.01TB", 0.01 * BYTES_PER_TB / BYTES_PER_MIB),
+        # Abbreviated forms
+        ("1K", BYTES_PER_KB / BYTES_PER_MIB),
+        ("10M", 10 * BYTES_PER_MB / BYTES_PER_MIB),
+        ("0.5g", 0.5 * BYTES_PER_GB / BYTES_PER_MIB),
+        ("0.01T", 0.01 * BYTES_PER_TB / BYTES_PER_MIB),
+        # With spaces
+        ("1 kiB", BYTES_PER_KIB / BYTES_PER_MIB),
+        ("10 MB", 10 * BYTES_PER_MB / BYTES_PER_MIB),
+        (" 0.5 GiB ", 0.5 * BYTES_PER_GIB / BYTES_PER_MIB),
+        ("0.01  TB", 0.01 * BYTES_PER_TB / BYTES_PER_MIB),
+    ],
+)
+def test_string_values_with_units(input_string, expected):
+    """Test various string inputs with different units"""
+    result = parse_size_to_megabytes(input_string, default_number_unit="GiB")
+    assert pytest.approx(result) == expected
+
+
+@pytest.mark.parametrize(
+    "invalid_input, default_unit",
+    [
+        ("invalid", "GiB"),  # Non-numeric input
+        ("1.2.3GiB", "GiB"),  # Invalid number format
+        ("1ZiB", "GiB"),  # Invalid unit
+        ("GiB", "GiB"),  # Missing number
+        ("1G1B", "GiB"),  # Invalid format
+        (5, None),  # Numeric input with default_number_unit as None
+        ("5", None),  # String numeric input with default_number_unit as None
+        ("5", "Gibb"),  # String numeric input with bad default_number_unit
+    ],
+)
+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)