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

maxframe/serialization/pandas.py

@@ -20,6 +20,11 @@ import pandas as pd
 from pandas.api.extensions import ExtensionArray
 from pandas.arrays import IntervalArray
 
+try:
+    from pandas.tseries.offsets import BaseOffset as PdBaseOffset
+except ImportError:
+    PdBaseOffset = type("FakeBaseOffset", (), {})
+
 from ..utils import no_default
 from .core import Serializer, buffered
 
@@ -39,14 +44,16 @@ class DataFrameSerializer(Serializer):
         self, serialized: List, context: Dict, subs: List[Any]
     ) -> pd.DataFrame:
         dtypes, idx = subs[:2]
-        seriess = [pd.Series(d, index=idx) for d in subs[2:]]
+        seriess = [
+            pd.Series(d, name=nm, index=idx).astype(dt)
+            for d, (nm, dt) in zip(subs[2:], dtypes.items())
+        ]
         if seriess:
-            df = pd.concat([pd.Series(d, index=idx) for d in subs[2:]], axis=1)
+            df = pd.concat(seriess, axis=1)
+            df.columns = dtypes.index
         else:
             df = pd.DataFrame([], index=idx)
-        df.columns = dtypes.index
-        df.index = idx
-        return df.astype(dtypes)
+        return df
 
 
 class SeriesSerializer(Serializer):
@@ -63,9 +70,26 @@ class SeriesSerializer(Serializer):
         return pd.Series(data, index=idx, name=name).astype(dtype)
 
 
+class PeriodSerializer(Serializer):
+    def serial(self, obj: pd.Period, context: Dict):
+        return [obj.strftime("%Y-%m-%d %H:%M:%S"), obj.freqstr], [], True
+
+    def deserial(self, serialized, context: Dict, subs: List):
+        return pd.Period(serialized[0], freq=serialized[1])
+
+
+class PdOffsetSerializer(Serializer):
+    def serial(self, obj: PdBaseOffset, context: Dict):
+        return [obj.freqstr], [], True
+
+    def deserial(self, serialized, context: Dict, subs: List):
+        return pd.tseries.frequencies.to_offset(serialized[0])
+
+
 _TYPE_CHAR_MULTI_INDEX = "M"
 _TYPE_CHAR_RANGE_INDEX = "R"
 _TYPE_CHAR_CATEGORICAL_INDEX = "C"
+_TYPE_CHAR_DATETIME_INDEX = "D"
 
 
 class IndexSerializer(Serializer):
@@ -80,6 +104,9 @@ class IndexSerializer(Serializer):
         elif isinstance(obj, pd.CategoricalIndex):
             data = [obj.name, obj.values]
             header = [_TYPE_CHAR_CATEGORICAL_INDEX]
+        elif isinstance(obj, pd.DatetimeIndex):
+            data = [obj.name, obj.values]
+            header = [_TYPE_CHAR_DATETIME_INDEX, obj.freqstr, None]
         else:
             if getattr(obj.dtype, "hasobject", False):
                 values = obj.tolist()
@@ -100,6 +127,10 @@ class IndexSerializer(Serializer):
         elif header[0] == _TYPE_CHAR_CATEGORICAL_INDEX:
             name, data = subs[:2]
             return pd.CategoricalIndex(data, name=name)
+        elif header[0] == _TYPE_CHAR_DATETIME_INDEX:
+            name, data = subs[:2]
+            freq, tz = header[1:]
+            return pd.DatetimeIndex(data, name=name, freq=freq, tz=tz)
         elif header[0] is None:  # Normal index
             dtype, name, values = subs
             return pd.Index(values, dtype=dtype, name=name)
@@ -162,9 +193,12 @@ class PdTimestampSerializer(Serializer):
             zone_info = []
             ts = obj.to_pydatetime().timestamp()
         elements = [int(ts), obj.microsecond, obj.nanosecond]
-        if hasattr(obj, "unit"):
-            elements.append(str(obj.unit))
-        return elements, zone_info, bool(zone_info)
+        for attr in ("unit", "freqstr"):
+            if getattr(obj, attr, None):
+                elements.append(str(getattr(obj, attr)))
+            else:
+                elements.append(None)
+        return elements, zone_info, not bool(zone_info)
 
     def deserial(self, serialized: List, context: Dict, subs: List):
         if subs:
@@ -196,7 +230,8 @@ class PdTimestampSerializer(Serializer):
                 "nanosecond": serialized[2],
             }
             if len(serialized) >= 4:
-                kwargs["unit"] = serialized[3]
+                ext_kw = dict(zip(("unit", "freq"), serialized[3:]))
+                kwargs.update({k: v for k, v in ext_kw.items() if v})
             val = pd.Timestamp(**kwargs)
         return val
 
@@ -238,4 +273,6 @@ CategoricalSerializer.register(pd.Categorical)
 ArraySerializer.register(ExtensionArray)
 PdTimestampSerializer.register(pd.Timestamp)
 PdTimedeltaSerializer.register(pd.Timedelta)
+PeriodSerializer.register(pd.Period)
+PdOffsetSerializer.register(PdBaseOffset)
 NoDefaultSerializer.register(type(no_default))

maxframe/serialization/serializables/core.py

@@ -453,8 +453,8 @@ class SerializableSerializer(Serializer):
             self._set_field_values(obj, primitives, cls_to_prim_key, True)
         if obj_class._NON_PRIMITIVE_FIELDS:
             self._set_field_values(obj, subs[0], cls_to_non_prim_key, False)
-        obj.__on_deserialize__()
-        return obj
+        ret = obj.__on_deserialize__()
+        return obj if ret is None else ret
 
 
 class NoFieldValueSerializer(Serializer):

maxframe/serialization/tests/test_serial.py

@@ -81,6 +81,11 @@ class CustomNamedTuple(NamedTuple):
         b"abcd",
         "abcd",
         slice(3, 9, 2),
+        slice(3, None, None),
+        slice(pd.Timestamp.now(), None, None),
+        range(10),
+        range(2, 10),
+        range(2, 10, 3),
         ["uvw", ("mno", "sdaf"), 4, 6.7],
         CustomNamedTuple("abcd", 13451),
         datetime.datetime.now(),
@@ -173,7 +178,7 @@ def test_timezones(val):
     ],
 )
 @switch_unpickle
-def test_numpy(val):
+def test_numpy_arrays(val):
     deserialized = deserialize(*serialize(val))
     assert type(val) == type(deserialized)
     np.testing.assert_equal(val, deserialized)
@@ -181,6 +186,13 @@ def test_numpy(val):
         assert deserialized.flags.f_contiguous
 
 
+@switch_unpickle
+def test_numpy_random_state():
+    rs1 = np.random.RandomState(131)
+    rs2 = deserialize(*serialize(rs1))
+    assert rs1.rand() == rs2.rand()
+
+
 @switch_unpickle
 def test_pandas():
     val = pd.Series([1, 2, 3, 4], name="nm")
@@ -213,16 +225,25 @@ def test_pandas():
     val = pd.CategoricalIndex(np.random.choice(list("abcd"), size=(1000,)))
     pd.testing.assert_index_equal(val, deserialize(*serialize(val)))
 
+    val = pd.DatetimeIndex([pd.Timestamp("2025-08-01 12:31:21")], freq="D")
+    pd.testing.assert_index_equal(val, deserialize(*serialize(val)))
+
+    val = pd.Period("2025-08-01", freq="D")
+    assert val == deserialize(*serialize(val))
+
+    val = pd.tseries.offsets.MonthEnd()
+    assert val == deserialize(*serialize(val))
+
 
 @switch_unpickle
 @pytest.mark.skipif(_arrow_dtype_supported, reason="pandas doesn't support ArrowDtype")
 def test_fake_arrow_dtype_serde():
     serializer = DtypeSerializer()
-    payload, data, ok = serializer.serial(
+    payload, data, is_leaf = serializer.serial(
         FakeArrowDtype(pa.map_(pa.int64(), pa.string())), dict()
     )
 
-    assert ok
+    assert is_leaf
     assert data == []
     assert payload == ["PA", "map<int64, string>"]
     new_dtype = serializer.deserial(payload, dict(), list())
@@ -233,6 +254,10 @@ def test_fake_arrow_dtype_serde():
 @switch_unpickle
 def test_arrow():
     test_array = np.random.rand(1000)
+    test_cplx_array = pa.array(
+        [["abc", "def"], ["ghi"], None],
+        type=pa.list_(pa.string()),
+    )
     test_df = pd.DataFrame(
         {
             "a": np.random.rand(1000),
@@ -242,14 +267,16 @@ def test_arrow():
     )
     test_vals = [
         pa.array(test_array),
+        test_cplx_array,
         pa.chunked_array([pa.array(test_array), pa.array(test_array)]),
+        pa.chunked_array([test_cplx_array, test_cplx_array]),
         pa.RecordBatch.from_pandas(test_df),
         pa.Table.from_pandas(test_df),
     ]
     for val in test_vals:
         deserialized = deserialize(*serialize(val))
         assert type(val) is type(deserialized)
-        np.testing.assert_equal(val, deserialized)
+        assert str(val) == str(deserialized)
 
 
 @pytest.mark.parametrize(

maxframe/tensor/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import extensions, fetch, linalg, random, special, statistics
+from . import extensions, fetch, fft, linalg, random, spatial, special, statistics
 from .arithmetic import absolute
 from .arithmetic import absolute as abs
 from .arithmetic import (
@@ -58,6 +58,7 @@ from .arithmetic import invert as bitwise_not
 from .arithmetic import (
     isclose,
     iscomplex,
+    iscomplexobj,
     isfinite,
     isinf,
     isnan,
@@ -152,32 +153,49 @@ from .indexing import (
     take,
     unravel_index,
 )
-from .linalg.dot import dot
-from .linalg.inner import inner, innerproduct
-from .linalg.matmul import matmul
-from .linalg.tensordot import tensordot
-from .linalg.vdot import vdot
+from .linalg import dot, einsum, inner, innerproduct, matmul, tensordot, vdot
 from .merge import append, column_stack, concatenate, dstack, hstack, stack, vstack
 from .misc import (
+    argwhere,
+    array_split,
     astype,
     atleast_1d,
     atleast_2d,
     atleast_3d,
+    broadcast_arrays,
     broadcast_to,
     copy,
+    copyto,
+    delete,
     diff,
+    dsplit,
+    ediff1d,
+    expand_dims,
+    flip,
+    fliplr,
+    flipud,
+    hsplit,
     in1d,
+    insert,
     isin,
+    moveaxis,
     ndim,
     ravel,
     repeat,
+    result_type,
+    roll,
+    rollaxis,
     searchsorted,
     setdiff1d,
+    shape,
+    split,
     squeeze,
     swapaxes,
+    tile,
     transpose,
     trapezoid,
     unique,
+    vsplit,
     where,
 )
 from .rechunk import rechunk
@@ -214,8 +232,20 @@ from .reduction import prod
 from .reduction import prod as product
 from .reduction import std, sum, var
 from .reshape import reshape
-from .sort import argsort, sort
-from .statistics import average, bincount, percentile, quantile
+from .sort import argpartition, argsort, partition, sort
+from .statistics import (
+    average,
+    bincount,
+    corrcoef,
+    cov,
+    digitize,
+    histogram,
+    histogram_bin_edges,
+    median,
+    percentile,
+    ptp,
+    quantile,
+)
 from .ufunc import ufunc
 
 # isort: off

maxframe/tensor/arithmetic/__init__.py

@@ -63,6 +63,7 @@ from .imag import TensorImag, imag
 from .invert import TensorInvert, invert
 from .isclose import TensorIsclose, isclose
 from .iscomplex import TensorIsComplex, iscomplex
+from .iscomplexobj import iscomplexobj
 from .isfinite import TensorIsFinite, isfinite
 from .isinf import TensorIsInf, isinf
 from .isnan import TensorIsNan, isnan
@@ -117,14 +118,6 @@ from .truediv import TensorTrueDiv, truediv
 from .trunc import TensorTrunc, trunc
 
 
-def _wrap_iop(func):
-    def inner(self, *args, **kwargs):
-        kwargs["out"] = self
-        return func(self, *args, **kwargs)
-
-    return inner
-
-
 def _install():
     from ..core import TENSOR_TYPE, Tensor, TensorData
     from ..datasource import tensor as astensor
@@ -150,6 +143,22 @@ def _install():
 
         return eq
 
+    def _wrap_iop(func):
+        def inner(self, *args, **kwargs):
+            kwargs["out"] = self
+            return func(self, *args, **kwargs)
+
+        return inner
+
+    def _wrap_magic(func):
+        def inner(lhs, rhs, **kwargs):
+            ret = func(lhs, rhs, **kwargs)
+            if isinstance(ret, TENSOR_TYPE):
+                ret.op.extra_params["magic"] = True
+            return ret
+
+        return inner
+
     for cls in TENSOR_TYPE:
         setattr(cls, "__add__", add)
         setattr(cls, "__iadd__", _wrap_iop(add))
@@ -182,8 +191,8 @@ def _install():
         setattr(cls, "__irshift__", _wrap_iop(rshift))
         setattr(cls, "__rrshift__", rrshift)
 
-        setattr(cls, "__eq__", _wrap_equal(equal))
-        setattr(cls, "__ne__", not_equal)
+        setattr(cls, "__eq__", _wrap_magic(_wrap_equal(equal)))
+        setattr(cls, "__ne__", _wrap_magic(not_equal))
         setattr(cls, "__lt__", less)
         setattr(cls, "__le__", less_equal)
         setattr(cls, "__gt__", greater)

maxframe/tensor/arithmetic/core.py

@@ -415,8 +415,8 @@ class TensorOutBinOp(TensorOperator, TensorElementWiseWithInputs):
         dtype = [r.dtype for r in self._fun(np.empty(1, dtype=x.dtype))]
 
         out = out or (None, None)
-        out1 = out1 or out[0]
-        out2 = out2 or out[1]
+        out1 = out1 if out1 is not None else out[0]
+        out2 = out2 if out2 is not None else out[1]
         x, out1, out2, where = self._process_inputs(x, out1, out2, where)
         shape = x.shape
         order1 = self._calc_order(x, out1)

maxframe/tensor/arithmetic/iscomplexobj.py

@@ -0,0 +1,53 @@
+# 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
+
+
+def iscomplexobj(x):
+    """
+    Check for a complex type or an array of complex numbers.
+
+    The type of the input is checked, not the value. Even if the input
+    has an imaginary part equal to zero, `iscomplexobj` evaluates to True.
+
+    Parameters
+    ----------
+    x : any
+        The input can be of any type and shape.
+
+    Returns
+    -------
+    iscomplexobj : bool
+        The return value, True if `x` is of a complex type or has at least
+        one complex element.
+
+    See Also
+    --------
+    isrealobj, iscomplex
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> mt.iscomplexobj(1)
+    False
+    >>> mt.iscomplexobj(1+0j)
+    True
+    >>> mt.iscomplexobj([3, 1+0j, True])
+    True
+    """
+    from ..datasource.array import tensor as astensor
+
+    t = astensor(x)
+    return issubclass(t.dtype.type, np.complexfloating)

maxframe/tensor/arithmetic/tests/test_arithmetic.py

@@ -16,7 +16,6 @@ import numpy as np
 import pytest
 import scipy.sparse as sps
 
-from ....core import enter_mode
 from ....utils import collect_leaf_operators
 from ...core import SparseTensor, Tensor
 from ...datasource import array, empty, ones, tensor
@@ -391,14 +390,6 @@ def test_get_set_real():
         a.real = [2, 4]
 
 
-def test_build_mode():
-    t1 = ones((2, 3), chunk_size=2)
-    assert t1 == 2
-
-    with enter_mode(build=True):
-        assert t1 != 2
-
-
 def test_unary_op_func_name():
     # make sure all the unary op has defined the func name.
 
@@ -450,3 +441,9 @@ def test_tree_arithmetic():
     assert len(t.inputs) == 3
     assert len(t.inputs[0].inputs) == 4
     assert len(t.inputs[-1].inputs) == 2
+
+
+def test_iscomplexobj():
+    assert not iscomplexobj(1)
+    assert iscomplexobj(1 + 0j)
+    assert iscomplexobj([3, 1 + 0j, True])

maxframe/tensor/core.py

@@ -35,6 +35,7 @@ from ..serialization.serializables import (
     Serializable,
     StringField,
 )
+from ..utils import make_dtype
 from .utils import fetch_corner_data, get_chunk_slices
 
 logger = logging.getLogger(__name__)
@@ -73,7 +74,7 @@ class TensorData(HasShapeTileableData, _ExecuteAndFetchMixin):
         super().__init__(
             _op=op,
             _shape=shape,
-            _dtype=dtype,
+            _dtype=make_dtype(dtype),
             _order=order,
             **kw,
         )
@@ -250,6 +251,9 @@ class Tensor(HasShapeTileable):
     def __len__(self):
         return len(self._data)
 
+    def __bool__(self):
+        return True if is_build_mode() else bool(self.to_numpy())
+
     @property
     def shape(self):
         return self._data.shape
@@ -520,7 +524,7 @@ class Tensor(HasShapeTileable):
         >>> a.execute()
         array([1, 2, 3, 4])
         """
-        from .misc import partition
+        from .sort import partition
 
         self._data = partition(self, kth, axis=axis, kind=kind, order=order, **kw).data
 

maxframe/tensor/datasource/tests/test_datasource.py

@@ -207,7 +207,8 @@ def test_set_tensor_inputs():
     assert t3.op.inputs == [t1.data, t2.data]
     assert t3.inputs == [t1.data, t2.data]
 
-    with pytest.raises(StopIteration):
+    msg = "Unexpected StopIteration happened"
+    with pytest.raises(RuntimeError, match=msg):
         t3.inputs = []
 
     t1 = tensor([1, 2], chunk_size=2)

maxframe/tensor/extensions/__init__.py

@@ -15,12 +15,14 @@
 from ...core import CachedAccessor
 from .accessor import TensorMaxFrameAccessor
 from .apply_chunk import apply_chunk
+from .rebalance import rebalance
 
 
 def _install():
     from ..core import TENSOR_TYPE
 
     TensorMaxFrameAccessor._register("apply_chunk", apply_chunk)
+    TensorMaxFrameAccessor._register("rebalance", rebalance)
 
     if TensorMaxFrameAccessor._api_count:
         for t in TENSOR_TYPE:

maxframe/tensor/extensions/apply_chunk.py

@@ -33,9 +33,9 @@ class TensorApplyChunk(TensorOperator, TensorOperatorMixin):
 
     func = FunctionField("func")
     elementwise = BoolField("elementwise")
-    args = TupleField("args")
-    kwargs = DictField("kwargs")
-    with_chunk_index = BoolField("with_chunk_index")
+    args = TupleField("args", default=None)
+    kwargs = DictField("kwargs", default=None)
+    with_chunk_index = BoolField("with_chunk_index", default=False)
 
     @classmethod
     def _set_inputs(cls, op: "TensorApplyChunk", inputs):

maxframe/tensor/extensions/rebalance.py

@@ -0,0 +1,65 @@
+# 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 ... import opcodes
+from ...serialization.serializables import Float64Field, Int64Field
+from ..operators import TensorOperator, TensorOperatorMixin
+
+
+class RebalanceMixin:
+    def __call__(self, in_obj):
+        self._output_types = in_obj.op.output_types
+        return self.new_tileable([in_obj], kws=[in_obj.params])
+
+
+class TensorRebalance(RebalanceMixin, TensorOperatorMixin, TensorOperator):
+    _op_type_ = opcodes.REBALANCE
+
+    axis = Int64Field("axis", default=None)
+    factor = Float64Field("factor", default=None)
+    num_partitions = Int64Field("num_partitions", default=None)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+
+def rebalance(tensor, axis=0, factor=None, num_partitions=None):
+    """
+    Make data more balanced across entire cluster.
+
+    Parameters
+    ----------
+    axis : int
+        The axis to rebalance.
+    factor : float
+        Specified so that number of chunks after balance is
+        total number of input chunks * factor.
+    num_partitions : int
+        Specified so the number of chunks are at most
+        num_partitions.
+
+    Returns
+    -------
+    Tensor
+        Result of tensor after rebalanced.
+    """
+    if num_partitions is None and factor is None:
+        raise ValueError("Need to specify num_partitions or factor")
+    if num_partitions is not None and factor is not None:
+        raise ValueError(
+            "num_partitions and factor cannot be specified at the same time"
+        )
+
+    op = TensorRebalance(axis=axis, factor=factor, num_partitions=num_partitions)
+    return op(tensor)

maxframe/tensor/fft/__init__.py

@@ -0,0 +1,32 @@
+# 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 .fft import TensorFFT, fft
+from .fft2 import TensorFFT2, fft2
+from .fftfreq import TensorFFTFreq, fftfreq
+from .fftn import TensorFFTN, fftn
+from .fftshift import TensorFFTShift, fftshift
+from .hfft import TensorHFFT, hfft
+from .ifft import TensorIFFT, ifft
+from .ifft2 import TensorIFFT2, ifft2
+from .ifftn import TensorIFFTN, ifftn
+from .ifftshift import TensorIFFTShift, ifftshift
+from .ihfft import TensorIHFFT, ihfft
+from .irfft import TensorIRFFT, irfft
+from .irfft2 import TensorIRFFT2, irfft2
+from .irfftn import TensorIRFFTN, irfftn
+from .rfft import TensorRFFT, rfft
+from .rfft2 import TensorRFFT2, rfft2
+from .rfftfreq import TensorRFFTFreq, rfftfreq
+from .rfftn import TensorRFFTN, rfftn