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

maxframe/io/odpsio/tests/test_schema.py

@@ -21,13 +21,15 @@ from odps import types as odps_types
 from .... import dataframe as md
 from .... import tensor as mt
 from ....core import OutputType
-from ....lib.dtypes_extension import dict_
+from ....lib.dtypes_extension import ArrowDtype, dict_, list_
 from ....utils import pd_release_version
 from ..schema import (
     arrow_schema_to_odps_schema,
     build_dataframe_table_meta,
     build_table_column_name,
     odps_schema_to_arrow_schema,
+    odps_schema_to_pandas_dtypes,
+    pandas_dtypes_to_arrow_schema,
     pandas_to_odps_schema,
     pandas_types_to_arrow_schema,
 )
@@ -42,6 +44,8 @@ def _wrap_maxframe_obj(obj, wrap="no"):
         obj = md.Series(obj)
     elif isinstance(obj, pd.Index):
         obj = md.Index(obj)
+    elif isinstance(obj, np.ndarray):
+        obj = mt.tensor(obj)
     else:
         obj = mt.scalar(obj)
     if wrap == "data":
@@ -191,6 +195,24 @@ def test_pandas_to_odps_schema_scalar(wrap_obj):
     assert meta.pd_index_level_names == [None]
 
 
+@pytest.mark.parametrize("wrap_obj", ["no", "yes", "data"])
+def test_pandas_to_odps_schema_tensor(wrap_obj):
+    data = np.array([1, 2, 3])
+
+    test_tensor = _wrap_maxframe_obj(data, wrap=wrap_obj)
+    if wrap_obj != "no":
+        test_tensor.op.data = None
+
+    schema, meta = pandas_to_odps_schema(test_tensor, unknown_as_string=True)
+    assert schema.columns[0].name == "_idx_0"
+    assert schema.columns[0].type.name == "bigint"
+    assert meta.type == OutputType.tensor
+    assert meta.table_column_names == []
+    assert meta.table_index_column_names == ["_idx_0"]
+    assert meta.pd_column_level_names == []
+    assert meta.pd_index_level_names == [None]
+
+
 def test_odps_arrow_schema_conversion():
     odps_schema = odps_types.OdpsSchema(
         [
@@ -275,6 +297,95 @@ def test_odps_arrow_schema_conversion():
         arrow_schema_to_odps_schema(pa.schema([("col1", pa.float16())]))
 
 
+def test_odps_pandas_schema_conversion():
+    odps_schema = odps_types.OdpsSchema(
+        [
+            odps_types.Column("col1", "string"),
+            odps_types.Column("col2", "binary"),
+            odps_types.Column("col3", "tinyint"),
+            odps_types.Column("col4", "smallint"),
+            odps_types.Column("col5", "int"),
+            odps_types.Column("col6", "bigint"),
+            odps_types.Column("col7", "boolean"),
+            odps_types.Column("col8", "float"),
+            odps_types.Column("col9", "double"),
+            # odps_types.Column("col10", "date"),
+            odps_types.Column("col11", "datetime"),
+            odps_types.Column("col12", "timestamp"),
+            # odps_types.Column("col13", "decimal(10, 2)"),
+            odps_types.Column("col14", "array<string>"),
+            odps_types.Column("col15", "map<string, bigint>"),
+            # odps_types.Column("col16", "struct<a1: string, a2: map<string, bigint>>"),
+            # odps_types.Column("col17", "CHAR(15)"),
+            # odps_types.Column("col18", "VARCHAR(15)"),
+            # odps_types.Column("col19", "decimal"),
+        ]
+    )
+    pd_dtypes = odps_schema_to_pandas_dtypes(odps_schema)
+    pd.testing.assert_series_equal(
+        pd_dtypes,
+        pd.Series(
+            [
+                np.dtype("O"),  # string
+                np.dtype("O"),  # binary
+                np.dtype(np.int8),
+                np.dtype(np.int16),
+                np.dtype(np.int32),
+                np.dtype(np.int64),
+                np.dtype(np.bool_),
+                np.dtype(np.float32),
+                np.dtype(np.float64),
+                np.dtype(
+                    "datetime64[ms]" if pd_release_version[0] >= 2 else "datetime64[ns]"
+                ),
+                np.dtype("datetime64[ns]"),
+                ArrowDtype(pa.list_(pa.string())),
+                ArrowDtype(pa.map_(pa.string(), pa.int64())),
+            ],
+            index=[c.name for c in odps_schema.columns],
+        ),
+    )
+
+    expected_odps_schema = odps_types.OdpsSchema(
+        [
+            odps_types.Column("col1", "string"),
+            odps_types.Column("col2", "string"),  # binary
+            odps_types.Column("col3", "tinyint"),
+            odps_types.Column("col4", "smallint"),
+            odps_types.Column("col5", "int"),
+            odps_types.Column("col6", "bigint"),
+            odps_types.Column("col7", "boolean"),
+            odps_types.Column("col8", "float"),
+            odps_types.Column("col9", "double"),
+            # odps_types.Column("col10", "date"),
+            odps_types.Column(
+                "col11", "datetime" if pd_release_version[0] >= 2 else "timestamp"
+            ),
+            odps_types.Column("col12", "timestamp"),
+            # odps_types.Column("col13", "decimal(10, 2)"),
+            odps_types.Column("col14", "array<string>"),
+            odps_types.Column("col15", "map<string, bigint>"),
+            # odps_types.Column("col16", "struct<a1: string, a2: map<string, bigint>>"),
+            # odps_types.Column("col17", "string"),
+            # odps_types.Column("col18", "string"),
+            # odps_types.Column("col19", "decimal(38, 18)"),
+        ]
+    )
+
+    odps_schema2 = arrow_schema_to_odps_schema(
+        pandas_dtypes_to_arrow_schema(pd_dtypes, unknown_as_string=True)
+    )
+    assert [c.name for c in expected_odps_schema.columns] == [
+        c.name for c in odps_schema2.columns
+    ]
+    assert [c.type for c in expected_odps_schema.columns] == [
+        c.type for c in odps_schema2.columns
+    ]
+
+    with pytest.raises(TypeError):
+        arrow_schema_to_odps_schema(pa.schema([("col1", pa.float16())]))
+
+
 def test_build_column_name():
     records = dict()
     assert build_table_column_name(0, "a" * 127, records) == "a" * 127
@@ -345,8 +456,10 @@ def test_pandas_types_to_arrow_schema():
         {
             "int8": pd.Series([], dtype=np.int8),
             "map": pd.Series([], dtype=dict_(pa.string(), pa.string())),
+            "list": pd.Series([], dtype=list_(pa.string())),
         },
     )
     schema = pandas_types_to_arrow_schema(pd_data)
     assert schema.field("int8").type == pa.int8()
     assert schema.field("map").type == pa.map_(pa.string(), pa.string())
+    assert schema.field("list").type == pa.list_(pa.string())

maxframe/io/odpsio/tests/test_tableio.py

@@ -14,15 +14,19 @@
 
 import datetime
 
+import mock
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pytest
 from odps import ODPS
+from odps.errors import TableModified
+from odps.models import Table
 
 from ....config import options
 from ....tests.utils import flaky, tn
 from ....utils import config_odps_default_options
+from .. import TunnelTableIO
 from ..tableio import ODPSTableIO
 
 
@@ -161,3 +165,41 @@ def test_table_io_with_parts(switch_table_io):
             pd.testing.assert_frame_equal(reader.read_all().to_pandas(), expected_data)
     finally:
         tb.drop()
+
+
+def test_tunnel_table_io_with_modified():
+    config_odps_default_options()
+
+    o = ODPS.from_environments()
+    table_io = TunnelTableIO(o)
+
+    # test read and write tables with partition
+    parted_table_name = tn("test_tunnel_write_modified")
+    o.delete_table(parted_table_name, if_exists=True)
+    tb = o.create_table(
+        parted_table_name,
+        (",".join(f"{c} double" for c in "abcde"), "pt string"),
+        lifecycle=1,
+    )
+
+    raised = False
+    raw_open_reader = Table.open_reader
+
+    def _new_open_reader(self, *args, **kwargs):
+        nonlocal raised
+        if not raised:
+            raised = True
+            raise TableModified("Intentional error")
+        return raw_open_reader(self, *args, **kwargs)
+
+    try:
+        pd_data = pd.DataFrame(np.random.rand(100, 5), columns=list("abcde"))
+        with table_io.open_writer(parted_table_name, "pt=test") as writer:
+            writer.write(pa.Table.from_pandas(pd_data, preserve_index=False))
+
+        with mock.patch(
+            "odps.models.table.Table.open_reader", new=_new_open_reader
+        ), table_io.open_reader(parted_table_name, "pt=test") as reader:
+            pd.testing.assert_frame_equal(reader.read_all().to_pandas(), pd_data)
+    finally:
+        tb.drop()

maxframe/io/odpsio/tests/test_volumeio.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
+
 import pytest
 from odps import ODPS
 
-from ....tests.utils import tn
+from ....tests.utils import create_test_volume, tn
 from ..volumeio import ODPSVolumeReader, ODPSVolumeWriter
 
 
@@ -24,59 +26,31 @@ def create_volume(request, oss_config):
     test_vol_name = tn("test_vol_name_" + request.param)
     odps_entry = ODPS.from_environments()
 
-    try:
-        odps_entry.delete_volume(test_vol_name)
-    except:
-        pass
+    @contextlib.contextmanager
+    def create_parted_volume():
+        try:
+            odps_entry.delete_volume(test_vol_name)
+        except:
+            pass
+        try:
+            odps_entry.create_parted_volume(test_vol_name)
+            yield
+        finally:
+            try:
+                odps_entry.delete_volume(test_vol_name)
+            except BaseException:
+                pass
 
     oss_test_dir_name = None
     if request.param == "parted":
-        odps_entry.create_parted_volume(test_vol_name)
+        ctx = create_parted_volume()
     else:
-        oss_test_dir_name = tn("test_oss_directory")
-        if oss_config is None:
-            pytest.skip("Need oss and its config to run this test")
-        (
-            oss_access_id,
-            oss_secret_access_key,
-            oss_bucket_name,
-            oss_endpoint,
-        ) = oss_config.oss_config
-
-        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, rolearn=rolearn
-        )
+        ctx = create_test_volume(test_vol_name, oss_config)
+
     try:
-        yield test_vol_name
+        with ctx:
+            yield test_vol_name
     finally:
-        try:
-            odps_entry.delete_volume(test_vol_name)
-        except BaseException:
-            pass
-
         if oss_test_dir_name is not None:
             import oss2
 

maxframe/learn/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import contrib
+from . import contrib, model_selection, preprocessing
 
-del contrib
+del contrib, model_selection, preprocessing

maxframe/learn/contrib/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import graph, llm, models, pytorch
+from . import graph, lightgbm, llm, pytorch
 
-del graph, llm, models, pytorch
+del graph, lightgbm, llm, pytorch

maxframe/learn/contrib/graph/connected_components.py

@@ -19,8 +19,9 @@ from maxframe import opcodes
 
 from ....core import OutputType
 from ....dataframe.operators import DataFrameOperator, DataFrameOperatorMixin
-from ....dataframe.utils import make_dtypes, parse_index
+from ....dataframe.utils import parse_index
 from ....serialization.serializables import Int32Field, StringField
+from ....utils import make_dtypes
 
 
 class DataFrameConnectedComponentsOperator(DataFrameOperator, DataFrameOperatorMixin):

maxframe/learn/contrib/lightgbm/__init__.py

@@ -0,0 +1,33 @@
+# 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 ._predict import predict
+from ._train import train
+from .dataset import Dataset
+
+# isort: off
+from ..utils import config_mod_getattr as _config_mod_getattr
+
+_config_mod_getattr(
+    {
+        "LGBMClassifier": ".classifier.LGBMClassifier",
+        "LGBMRegressor": ".regressor.LGBMRegressor",
+        # "LGBMRanker": ".ranker.LGBMRanker",
+    },
+    globals(),
+)
+
+del _config_mod_getattr
+
+__all__ = ["Dataset", "LGBMClassifier", "LGBMRegressor", "predict", "train"]

maxframe/learn/contrib/lightgbm/_predict.py

@@ -0,0 +1,138 @@
+# 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 Any, Dict, List
+
+import numpy as np
+
+from .... import opcodes
+from ....core import OutputType
+from ....core.operator import Operator, TileableOperatorMixin
+from ....serialization.serializables import BoolField, DictField, Int32Field, KeyField
+from ....tensor.core import TensorOrder
+from ....typing_ import EntityType
+from ..models import to_remote_model
+from .dataset import check_data
+
+
+class LGBMPredict(Operator, TileableOperatorMixin):
+    _op_type_ = opcodes.LGBM_PREDICT
+    _op_module_ = "learn.contrib.lightgbm"
+    output_dtype = np.dtype(np.float64)
+
+    data = KeyField("data", default=None)
+    booster = KeyField("booster", default=None)
+    start_iteration = Int32Field("start_iteration", default=0)
+    num_iterations = Int32Field("num_iterations", default=None)
+    raw_score = BoolField("raw_score", default=False)
+    pred_leaf = BoolField("pred_leaf", default=False)
+    pred_contrib = BoolField("pred_contrib", default=False)
+    validate_features = BoolField("validate_features", default=False)
+    kwds = DictField("kwds", default=None)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+    def has_custom_code(self) -> bool:
+        return True
+
+    @property
+    def output_limit(self) -> int:
+        return 1 + self.pred_leaf + self.pred_contrib
+
+    @classmethod
+    def _set_inputs(cls, op: "LGBMPredict", inputs: List[EntityType]):
+        super()._set_inputs(op, inputs)
+        it = iter(inputs)
+        op.data = next(it)
+        op.booster = next(it)
+
+    def __call__(self):
+        num_class = getattr(self.booster.op, "num_class", None)
+
+        result_kw: Dict[str, Any] = {
+            "dtype": self.output_dtype,
+            "order": TensorOrder.C_ORDER,
+        }
+        kws = [result_kw]
+
+        if num_class is not None:
+            num_class = int(num_class)
+        if num_class is not None and num_class > 2:
+            result_kw["shape"] = (self.data.shape[0], num_class)
+        else:
+            result_kw["shape"] = (self.data.shape[0],)
+
+        if self.pred_leaf:
+            kws.append(
+                {
+                    "shape": (np.nan, np.nan),
+                    "dtype": np.dtype(np.float_),
+                    "order": TensorOrder.C_ORDER,
+                }
+            )
+        if self.pred_contrib:
+            kws.append(
+                {
+                    "shape": (np.nan, np.nan),
+                    "dtype": np.dtype(np.float_),
+                    "order": TensorOrder.C_ORDER,
+                }
+            )
+
+        return self.new_tileables([self.data, self.booster], kws=kws)
+
+
+def predict(
+    booster,
+    data,
+    raw_score: bool = False,
+    start_iteration: int = 0,
+    num_iteration: int = None,
+    pred_leaf: bool = False,
+    pred_contrib: bool = False,
+    validate_features: bool = False,
+    **kwargs,
+):
+    import lightgbm
+
+    from .core import Booster, BoosterData
+
+    if not isinstance(booster, (Booster, BoosterData, lightgbm.Booster)):
+        raise TypeError(
+            f"model has to be a lightgbm.Booster, got {type(booster)} instead"
+        )
+    elif isinstance(booster, lightgbm.Booster):
+        booster = to_remote_model(booster, model_cls=Booster)
+    proba = kwargs.pop("proba", False)
+
+    data = check_data(data)
+
+    op = LGBMPredict(
+        data=data,
+        booster=booster,
+        output_types=[OutputType.tensor],
+        proba=proba,
+        raw_score=raw_score,
+        start_iteration=start_iteration,
+        num_iteration=num_iteration,
+        pred_leaf=pred_leaf,
+        pred_contrib=pred_contrib,
+        validate_features=validate_features,
+        kwds=kwargs,
+    )
+    results = op()
+    if len(results) == 1:
+        return results[0]
+    return results

maxframe/learn/contrib/lightgbm/_train.py

@@ -0,0 +1,163 @@
+# 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 itertools
+import logging
+from typing import List
+
+from .... import opcodes
+from ....core import OutputType
+from ....core.operator import ObjectOperator, ObjectOperatorMixin
+from ....serialization.serializables import (
+    AnyField,
+    BoolField,
+    DictField,
+    FieldTypes,
+    FunctionField,
+    Int32Field,
+    KeyField,
+    ListField,
+    StringField,
+)
+from ....typing_ import EntityType
+from ..models import to_remote_model
+from ..utils import TrainingCallback
+from .core import Booster, BoosterData, LGBMScikitLearnBase
+
+logger = logging.getLogger(__name__)
+
+
+class LGBMTrain(ObjectOperator, ObjectOperatorMixin):
+    _op_type_ = opcodes.LGBM_TRAIN
+    _op_module_ = "learn.contrib.lightgbm"
+
+    params = DictField("params", key_type=FieldTypes.string, default_factory=dict)
+    train_set = KeyField("train_set", default=None)
+    num_boost_round = Int32Field("num_boost_round", default=None)
+    valid_sets = ListField("valid_sets", FieldTypes.key, default_factory=list)
+    valid_names = ListField("valid_names", FieldTypes.string, default_factory=list)
+    feval = FunctionField("feval", default=None)
+    init_model = KeyField("init_model", default=None)
+    feature_name = AnyField("feature_name", default=None)
+    categorical_feature = AnyField("categorical_feature", default=None)
+    keep_training_booster = BoolField("keep_training_booster", default=False)
+    callbacks = ListField(
+        "callbacks",
+        field_type=FunctionField.field_type,
+        default=None,
+        on_serialize=TrainingCallback.from_local,
+    )
+    tree_learner = StringField("tree_learner", default=None)
+    timeout = Int32Field("timeout", default=None)
+    # indicating shape of the predicted data of the model
+    num_class = Int32Field("num_class", default=None)
+
+    def __init__(self, gpu=None, **kw):
+        super().__init__(gpu=gpu, **kw)
+        if self.output_types is None:
+            self.output_types = [OutputType.object]
+        if self.has_evals_result:
+            self.output_types.append(OutputType.object)
+
+    @classmethod
+    def _set_inputs(cls, op: "LGBMTrain", inputs: List[EntityType]):
+        super()._set_inputs(op, inputs)
+        input_it = iter(op.inputs)
+        op.train_set = next(input_it)
+        op.valid_sets = list(itertools.islice(input_it, len(op.valid_sets or [])))
+        if op.init_model is not None:
+            op.init_model = next(input_it)
+
+    def __call__(self, evals_result):
+        inputs = [self.train_set]
+        if self.has_evals_result:
+            inputs.extend(self.valid_sets)
+        if self.init_model is not None:
+            inputs.append(self.init_model)
+        kws = [{"object_class": Booster}, {}]
+        return self.new_tileables(inputs, kws=kws, evals_result=evals_result)[0]
+
+    @property
+    def output_limit(self):
+        return 2 if self.has_evals_result else 1
+
+    def has_custom_code(self) -> bool:
+        if not self.callbacks:
+            return False
+        return any(
+            not isinstance(cb, TrainingCallback) or cb.has_custom_code()
+            for cb in self.callbacks
+        )
+
+    @property
+    def has_evals_result(self) -> bool:
+        return bool(self.valid_sets)
+
+
+def _get_lgbm_booster(init_model):
+    import lightgbm
+
+    if isinstance(init_model, (LGBMScikitLearnBase, lightgbm.LGBMModel)):
+        init_model = init_model.booster_
+
+    if isinstance(init_model, (Booster, BoosterData)):
+        return init_model
+    elif isinstance(init_model, lightgbm.Booster):
+        return to_remote_model(init_model, model_cls=Booster)
+    raise ValueError(f"Cannot use {type(init_model)} as init_model")
+
+
+def train(
+    params,
+    train_set,
+    num_boost_round=100,
+    valid_sets=None,
+    valid_names=None,
+    feval=None,
+    init_model=None,
+    keep_training_booster=False,
+    callbacks=None,
+    num_class=2,
+    evals_result=None,
+    **kwargs,
+):
+    evals_result = evals_result if evals_result is not None else dict()
+    timeout = kwargs.pop("timeout", 120)
+    session = kwargs.pop("session", None)
+    run_kwargs = kwargs.pop("run_kwargs", dict())
+    if valid_sets and valid_names:
+        if not isinstance(valid_names, list):
+            raise TypeError("valid_names must be a list of strings")
+        if len(valid_names) != len(valid_sets):
+            raise ValueError("size of valid_names must match size of valid_sets")
+    if init_model is not None:
+        init_model = _get_lgbm_booster(init_model)
+    data = LGBMTrain(
+        params=params,
+        train_set=train_set,
+        num_boost_round=num_boost_round,
+        valid_sets=valid_sets,
+        valid_names=valid_names,
+        feval=feval,
+        init_model=init_model,
+        keep_training_booster=keep_training_booster,
+        callbacks=callbacks,
+        num_class=num_class,
+        evals_result=evals_result,
+        timeout=timeout,
+        **kwargs,
+    )(evals_result)
+    if valid_sets:
+        data.execute(session=session, **run_kwargs)
+    return data