maxframe 1.0.0rc4__cp310-cp310-win_amd64.whl → 1.1.1__cp310-cp310-win_amd64.whl

maxframe/dataframe/tests/test_utils.py

@@ -0,0 +1,60 @@
+# Copyright 1999-2024 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
+import pytest
+
+from ...udf import MarkedFunction, with_python_requirements, with_resources
+from ..utils import pack_func_args
+
+
+@pytest.fixture
+def df1():
+    return pd.DataFrame([[1, 2], [3, 4]], columns=["a", "b"])
+
+
+def test_pack_function(df1):
+    # pack normal function
+    @with_resources("a.zip")
+    def keep(df):
+        return df
+
+    f = pack_func_args(df1, keep)
+    assert f(df1).equals(df1)
+    assert isinstance(f, MarkedFunction)
+    assert f.resources == ["a.zip"]
+
+    # pack with args
+    @with_python_requirements("numpy")
+    def add(a, b):
+        return a + b
+
+    f = pack_func_args(df1, add, 1)
+    assert f(df1).equals(df1 + 1)
+    assert isinstance(f, MarkedFunction)
+    assert f.pythonpacks[0].requirements == ("numpy",)
+
+    f = pack_func_args(df1, np.sum)
+    assert f(df1).equals(np.sum(df1))
+
+    @with_resources("a.txt")
+    @with_python_requirements("pandas")
+    def times_add(df, param, times):
+        return df * times + param
+
+    f = pack_func_args(df1, times_add, 5, 6)
+    assert f(df1).equals(df1 * 6 + 5)
+    assert isinstance(f, MarkedFunction)
+    assert f.resources == ["a.txt"]
+    assert f.pythonpacks[0].requirements == ("pandas",)

maxframe/dataframe/utils.py

@@ -20,7 +20,7 @@ import operator
 import sys
 from contextlib import contextmanager
 from numbers import Integral
-from typing import Any, Callable, List
+from typing import TYPE_CHECKING, Any, Callable, List
 
 import numpy as np
 import pandas as pd
@@ -30,6 +30,7 @@ from pandas.core.dtypes.inference import is_dict_like, is_list_like
 
 from ..core import Entity, ExecutableTuple
 from ..lib.mmh3 import hash as mmh_hash
+from ..udf import MarkedFunction
 from ..utils import (
     ModulePlaceholder,
     is_full_slice,
@@ -44,6 +45,9 @@ try:
 except ImportError:  # pragma: no cover
     pa = ModulePlaceholder("pyarrow")
 
+if TYPE_CHECKING:
+    from .operators import DataFrameOperator
+
 cudf = lazy_import("cudf", rename="cudf")
 vineyard = lazy_import("vineyard")
 try:
@@ -263,12 +267,30 @@ def parse_index(index_value, *args, store_data=False, key=None):
         return IndexValue(_index_value=_serialize_index(index_value))
 
 
-def gen_unknown_index_value(index_value, *args):
+def gen_unknown_index_value(index_value, *args, normalize_range_index=False):
+    """
+    Generate new index value with the same likes of given index_value and args, but without any value.
+
+    Parameters
+    ----------
+    index_value
+        Given index value.
+    args
+        Arguments for parse_index.
+    normalize_range_index
+        If normalize range index to normal index.
+
+    Returns
+    -------
+        New created range index value.
+    """
     pd_index = index_value.to_pandas()
-    if isinstance(pd_index, pd.RangeIndex):
-        return parse_index(pd.RangeIndex(-1), *args)
+    if not normalize_range_index and isinstance(pd_index, pd.RangeIndex):
+        return parse_index(pd.RangeIndex(-1, name=pd_index.name), *args)
     elif not isinstance(pd_index, pd.MultiIndex):
-        return parse_index(pd.Index([], dtype=pd_index.dtype), *args)
+        return parse_index(
+            pd.Index([], dtype=pd_index.dtype, name=pd_index.name), *args
+        )
     else:
         i = pd.MultiIndex.from_arrays(
             [c[:0] for c in pd_index.levels], names=pd_index.names
@@ -1160,7 +1182,65 @@ def patch_sa_engine_execute():
     Engine.execute = execute
 
 
-def pack_func_args(df, funcs, *args, **kwargs) -> Any:
+def bind_func_args_from_pos(func, args_bind_position, *bound_args, **bound_kwargs):
+    """
+    Create a new function with arguments bound from specified position.
+
+    Parameters
+    ----------
+    func : callable
+        Target function to be wrapped.
+    args_bind_position : int
+        Position to start binding arguments (0-based).
+        e.g., n=0 binds from first arg, n=1 binds from second arg.
+    *bound_args : tuple
+        Arguments to be bound from position n.
+    **bound_kwargs : dict
+        Keyword arguments to be bound.
+
+    Returns
+    -------
+    callable
+        Wrapped function with bound arguments.
+
+    Examples
+    --------
+    >>> def func(x, y, z=0):
+    ...    return x * y + z
+    >>> f = bind_func_args_from_pos(func, 0, 10)  # bind from second position
+    >>> f(5)  # equals func(5, 10)
+    10
+
+    Raises
+    ------
+    TypeError
+        If func is not callable or n is not an integer.
+    ValueError
+        If n is negative or exceeds the number of parameters.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*runtime_args, **runtime_kwargs):
+        try:
+            # Combine arguments
+            all_args = (
+                runtime_args[:args_bind_position]
+                + bound_args
+                + runtime_args[args_bind_position:]
+            )
+            all_kwargs = {**bound_kwargs, **runtime_kwargs}
+
+            return func(*all_args, **all_kwargs)
+        except Exception as e:
+            # Enhance error message with context
+            raise type(e)(
+                f"Error calling {func.__name__} with bound arguments: {str(e)}"
+            ) from e
+
+    return wrapper
+
+
+def pack_func_args(df, funcs, *args, args_bind_position=1, **kwargs) -> Any:
     """
     Pack the funcs with args and kwargs to avoid the ambiguity between other
     positional and keyword arguments. It will process the funcs by the following rule:
@@ -1189,6 +1269,9 @@ def pack_func_args(df, funcs, *args, **kwargs) -> Any:
         The DataFrame or Series object to test the function.
     funcs : function, str, list-like or dict-like
         Function to pack. It should have the same type with Dataframe.transform().
+    args_bind_position: int
+        Position to start binding arguments (0-based).
+            e.g., n=0 binds from first arg, n=1 binds from second arg.
     *args :
         The positional arguments to func. If funcs contains many functions, each one
         should be able to accept *args.
@@ -1219,8 +1302,19 @@ def pack_func_args(df, funcs, *args, **kwargs) -> Any:
 
     f = get_callable_by_name(df, funcs) if isinstance(funcs, str) else funcs
 
+    from ..udf import MarkedFunction
+
+    if isinstance(f, MarkedFunction):
+        # for marked function, pack the inner function, and reset as mark function
+        packed_func = f.copy()
+        packed_func.func = bind_func_args_from_pos(
+            f.func, args_bind_position, *args, **kwargs
+        )
+    else:
+        packed_func = bind_func_args_from_pos(f, args_bind_position, *args, **kwargs)
+
     # Callable
-    return functools.partial(f, *args, **kwargs)
+    return packed_func
 
 
 def get_callable_by_name(df: Any, func_name: str) -> Callable:
@@ -1262,3 +1356,12 @@ def get_callable_by_name(df: Any, func_name: str) -> Callable:
     raise AttributeError(
         f"'{func_name}' is not a valid function for '{type(df).__name__}' object"
     )
+
+
+def copy_func_scheduling_hints(func, op: "DataFrameOperator") -> None:
+    if not isinstance(func, MarkedFunction):
+        return
+    if func.expect_engine:
+        op.expect_engine = func.expect_engine
+    if func.expect_resources:
+        op.expect_resources = func.expect_resources

maxframe/dataframe/window/expanding.py

@@ -28,6 +28,7 @@ from .aggregation import BaseDataFrameExpandingAgg
 from .core import Window
 
 _window_has_method = pd_release_version >= (1, 3, 0)
+_window_has_center = pd_release_version < (2, 0, 0)
 
 
 class DataFrameExpandingAgg(BaseDataFrameExpandingAgg):
@@ -49,10 +50,11 @@ class Expanding(Window):
     def params(self):
         p = OrderedDict()
 
+        args = ["min_periods", "center", "axis", "method"]
         if not _window_has_method:  # pragma: no cover
-            args = ["min_periods", "center", "axis"]
-        else:
-            args = ["min_periods", "center", "axis", "method"]
+            args = [a for a in args if a != "method"]
+        if not _window_has_center:
+            args = [a for a in args if a != "center"]
 
         for k in args:
             p[k] = getattr(self, k)

maxframe/dataframe/window/tests/test_expanding.py

@@ -29,8 +29,8 @@ def test_expanding():
     with pytest.raises(NotImplementedError):
         _ = df2.expanding(3, axis=1)
 
-    r = df2.expanding(3, center=False)
-    expected = df.expanding(3, center=False)
+    r = df2.expanding(3)
+    expected = df.expanding(3)
     assert repr(r) == repr(expected)
 
     assert "b" in dir(r)

maxframe/io/objects/tests/test_object_io.py

@@ -11,6 +11,7 @@
 # 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 pytest
 from odps import ODPS
@@ -48,15 +49,33 @@ def create_volume(request, oss_config):
         oss_bucket_name,
         oss_endpoint,
     ) = oss_config.oss_config
-    test_location = "oss://%s:%s@%s/%s/%s" % (
-        oss_access_id,
-        oss_secret_access_key,
-        oss_endpoint,
-        oss_bucket_name,
-        oss_test_dir_name,
-    )
+
+    if "test" in oss_endpoint:
+        # offline config
+        test_location = "oss://%s:%s@%s/%s/%s" % (
+            oss_access_id,
+            oss_secret_access_key,
+            oss_endpoint,
+            oss_bucket_name,
+            oss_test_dir_name,
+        )
+        rolearn = None
+    else:
+        # online config
+        endpoint_parts = oss_endpoint.split(".", 1)
+        if "-internal" not in endpoint_parts[0]:
+            endpoint_parts[0] += "-internal"
+        test_location = "oss://%s/%s/%s" % (
+            ".".join(endpoint_parts),
+            oss_bucket_name,
+            oss_test_dir_name,
+        )
+        rolearn = oss_config.oss_rolearn
+
     oss_config.oss_bucket.put_object(oss_test_dir_name + "/", b"")
-    odps_entry.create_external_volume(test_vol_name, location=test_location)
+    odps_entry.create_external_volume(
+        test_vol_name, location=test_location, rolearn=rolearn
+    )
 
     try:
         yield test_vol_name
@@ -75,8 +94,12 @@ def test_simple_object_io(create_volume):
 
     odps_entry = ODPS.from_environments()
 
-    reader = ODPSVolumeReader(odps_entry, create_volume, obj.key)
-    writer = ODPSVolumeWriter(odps_entry, create_volume, obj.key)
+    reader = ODPSVolumeReader(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
+    writer = ODPSVolumeWriter(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
 
     handler = get_object_io_handler(obj)()
     handler.write_object(writer, obj, data)
@@ -89,8 +112,12 @@ def test_tensor_object_io(create_volume):
 
     odps_entry = ODPS.from_environments()
 
-    reader = ODPSVolumeReader(odps_entry, create_volume, obj.key)
-    writer = ODPSVolumeWriter(odps_entry, create_volume, obj.key)
+    reader = ODPSVolumeReader(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
+    writer = ODPSVolumeWriter(
+        odps_entry, create_volume, obj.key, replace_internal_host=True
+    )
 
     handler = get_object_io_handler(obj)()
     handler.write_object(writer, obj, data)

maxframe/io/odpsio/__init__.py

@@ -19,5 +19,5 @@ from .schema import (
     odps_schema_to_pandas_dtypes,
     pandas_to_odps_schema,
 )
-from .tableio import HaloTableIO, ODPSTableIO
+from .tableio import HaloTableIO, ODPSTableIO, TunnelTableIO
 from .volumeio import ODPSVolumeReader, ODPSVolumeWriter

maxframe/io/odpsio/arrow.py

@@ -14,10 +14,12 @@
 
 from typing import Any, Tuple, Union
 
+import numpy as np
 import pandas as pd
 import pyarrow as pa
 
 from ...core import OutputType
+from ...lib.version import parse as parse_version
 from ...protocol import DataFrameTableMeta
 from ...tensor.core import TENSOR_TYPE
 from ...typing_ import ArrowTableType, PandasObjectTypes
@@ -69,13 +71,24 @@ def arrow_to_pandas(
 
 
 def pandas_to_arrow(
-    df: Any, nthreads=1, ignore_index=False
+    df: Any, nthreads=1, ignore_index=False, ms_cols=None
 ) -> Tuple[ArrowTableType, DataFrameTableMeta]:
     table_meta = build_dataframe_table_meta(df, ignore_index)
     df = df.copy() if callable(getattr(df, "copy", None)) else df
+    table_datetime_cols = None
     if table_meta.type in (OutputType.dataframe, OutputType.series):
         if table_meta.type == OutputType.series:
             df = df.to_frame("_data" if df.name is None else df.name)
+            if ms_cols:
+                table_datetime_cols = {"_data"}
+        elif ms_cols:
+            ms_col_set = set(ms_cols)
+            table_datetime_cols = set()
+            for pd_col, table_col in zip(
+                table_meta.pd_column_dtypes.keys(), table_meta.table_column_names
+            ):
+                if pd_col in ms_col_set:
+                    table_datetime_cols.add(table_col)
         df.columns = pd.Index(table_meta.table_column_names)
         if not ignore_index:
             df = df.rename_axis(table_meta.table_index_column_names).reset_index()
@@ -83,6 +96,12 @@ def pandas_to_arrow(
         df = pd.DataFrame([], columns=[])
     elif table_meta.type == OutputType.index:
         names = [f"_idx_{idx}" for idx in range(len(df.names))]
+        table_datetime_cols = set()
+        if ms_cols:
+            if isinstance(df, pd.MultiIndex):
+                table_datetime_cols = {f"_idx_{idx}" for idx in ms_cols}
+            else:
+                table_datetime_cols = {"_idx_0"}
         df = df.to_frame(name=names[0] if len(names) == 1 else names)
     elif table_meta.type == OutputType.scalar:
         names = ["_idx_0"]
@@ -92,4 +111,34 @@ def pandas_to_arrow(
             df = pd.DataFrame([[df]], columns=names)
     else:  # this could never happen  # pragma: no cover
         raise ValueError(f"Does not support meta type {table_meta.type!r}")
-    return pa.Table.from_pandas(df, nthreads=nthreads, preserve_index=False), table_meta
+
+    try:
+        pa_table = pa.Table.from_pandas(df, nthreads=nthreads, preserve_index=False)
+    except pa.ArrowTypeError as ex:  # pragma: no cover
+        late_np_version = parse_version(np.__version__) >= parse_version("1.20")
+        early_pa_version = parse_version(pa.__version__) <= parse_version("4.0")
+        if (
+            late_np_version
+            and early_pa_version
+            and "Did not pass numpy.dtype object" in str(ex)
+        ):
+            raise TypeError(
+                "Potential dependency conflict. Try update to pyarrow>4.0 "
+                "or downgrade to numpy<1.20. Details can be seen at "
+                "https://github.com/numpy/numpy/issues/17913. "
+                f"Raw error message: {ex!r}"
+            ).with_traceback(ex.__traceback__) from None
+        else:
+            raise
+
+    if table_datetime_cols:
+        col_names = pa_table.schema.names
+        col_datas = []
+        for idx, col_name in enumerate(pa_table.schema.names):
+            if col_name not in table_datetime_cols:
+                col_datas.append(pa_table.column(idx))
+                continue
+            col_data = pa_table.column(idx).cast(pa.timestamp("ms"))
+            col_datas.append(col_data)
+        pa_table = pa.Table.from_arrays(col_datas, names=col_names)
+    return pa_table, table_meta

maxframe/io/odpsio/schema.py

@@ -16,6 +16,7 @@ import string
 from collections import defaultdict
 from typing import Any, Dict, Tuple
 
+import numpy as np
 import pandas as pd
 import pyarrow as pa
 from odps import types as odps_types
@@ -39,6 +40,7 @@ _arrow_to_odps_types = {
     pa.float64(): odps_types.double,
     pa.date32(): odps_types.date,
     pa.timestamp("ms"): odps_types.datetime,
+    pa.timestamp("us"): odps_types.timestamp,
     pa.timestamp("ns"): odps_types.timestamp,
 }
 
@@ -205,20 +207,35 @@ def pandas_to_odps_schema(
     else:
         empty_columns = None
 
+    ms_cols = None
     if isinstance(df_obj, (md.DataFrame, pd.DataFrame)):
         empty_df_obj = pd.DataFrame(
             [], columns=empty_columns, index=empty_index
         ).astype(df_obj.dtypes)
+        ms_cols = [
+            col for col, dt in df_obj.dtypes.items() if dt == np.dtype("datetime64[ms]")
+        ]
     elif isinstance(df_obj, (md.Series, pd.Series)):
         empty_df_obj = pd.Series([], name=df_obj.name, index=empty_index).astype(
             df_obj.dtype
         )
+        ms_cols = df_obj.dtype == np.dtype("datetime64[ms]")
     elif isinstance(df_obj, (md.Index, pd.Index)):
         empty_df_obj = empty_index
+        if isinstance(empty_index, pd.MultiIndex):
+            ms_cols = [
+                idx
+                for idx, dt in enumerate(empty_index.dtypes.values)
+                if dt == np.dtype("datetime64[ms]")
+            ]
+        else:
+            ms_cols = df_obj.dtype == np.dtype("datetime64[ms]")
     else:
         empty_df_obj = df_obj
 
-    arrow_data, table_meta = pandas_to_arrow(empty_df_obj, ignore_index=ignore_index)
+    arrow_data, table_meta = pandas_to_arrow(
+        empty_df_obj, ignore_index=ignore_index, ms_cols=ms_cols
+    )
     return (
         arrow_schema_to_odps_schema(
             arrow_data.schema, unknown_as_string=unknown_as_string
@@ -346,10 +363,11 @@ def build_dataframe_table_meta(
     else:
         pd_index_val = index_obj
 
-    if hasattr(pd_index_val, "dtypes"):
-        index_dtypes = pd.Series(pd_index_val.dtypes.values, index=pd_index_val.names)
-    else:
-        index_dtypes = pd.Series([pd_index_val.dtype], index=pd_index_val.names)
+    level_dtypes = [
+        pd_index_val.get_level_values(level).dtype
+        for level in range(pd_index_val.nlevels)
+    ]
+    index_dtypes = pd.Series(level_dtypes, index=pd_index_val.names)
 
     if ignore_index and obj_type != OutputType.index:
         table_index_column_names = []