maxframe 1.1.0__cp311-cp311-macosx_10_9_universal2.whl → 1.2.0__cp311-cp311-macosx_10_9_universal2.whl

maxframe/io/odpsio/arrow.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.
@@ -14,14 +14,20 @@
 
 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
-from .schema import build_dataframe_table_meta
+from .schema import (
+    arrow_table_to_pandas_dataframe,
+    build_dataframe_table_meta,
+    pandas_dataframe_to_arrow_table,
+)
 
 
 def _rebuild_dataframe(
@@ -57,7 +63,7 @@ def _rebuild_index(df: pd.DataFrame, table_meta: DataFrameTableMeta) -> pd.Index
 def arrow_to_pandas(
     arrow_table: ArrowTableType, table_meta: DataFrameTableMeta
 ) -> PandasObjectTypes:
-    df = arrow_table.to_pandas()
+    df = arrow_table_to_pandas_dataframe(arrow_table, table_meta)
     if table_meta.type in (OutputType.dataframe, OutputType.series):
         return _rebuild_dataframe(df, table_meta)
     elif table_meta.type == OutputType.index:
@@ -109,7 +115,26 @@ 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}")
-    pa_table = pa.Table.from_pandas(df, nthreads=nthreads, preserve_index=False)
+
+    try:
+        pa_table = pandas_dataframe_to_arrow_table(df, nthreads=nthreads)
+    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 = []

maxframe/io/odpsio/schema.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.
@@ -14,7 +14,7 @@
 
 import string
 from collections import defaultdict
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -23,6 +23,7 @@ from odps import types as odps_types
 from pandas.api import types as pd_types
 
 from ...core import TILEABLE_TYPE, OutputType
+from ...lib.dtypes_extension import ArrowDtype
 from ...protocol import DataFrameTableMeta
 from ...tensor.core import TENSOR_TYPE
 
@@ -61,6 +62,33 @@ _odps_type_to_arrow = {
     odps_types.timestamp_ntz: pa.timestamp("ns"),
 }
 
+_based_for_pandas_pa_dtypes = Union[pa.MapType]
+
+
+def is_based_for_pandas_dtype(dtype: pa.DataType) -> bool:
+    """
+    Check whether the arrow type is based for one pandas data type.
+    If true, we should make sure the environment support ArrowDtype.
+    """
+    if not isinstance(dtype, _based_for_pandas_pa_dtypes):
+        return False
+
+    if ArrowDtype is None:
+        raise ImportError("ArrowDtype is not supported in current environment")
+    return True
+
+
+def pandas_types_to_arrow_schema(df_obj: pd.DataFrame) -> pa.Schema:
+    """
+    This one is only called when a pandas DataFrame is written to ODPS. So we can check
+    whether the ArrowDtype is supported.
+    """
+    schema = pa.Schema.from_pandas(df_obj, preserve_index=False)
+    for idx, col_dtype in enumerate(df_obj.dtypes.items()):
+        if ArrowDtype is not None and isinstance(col_dtype[1], ArrowDtype):
+            schema.set(idx, pa.field(col_dtype[0], col_dtype[1].pyarrow_dtype))
+    return schema
+
 
 def arrow_type_to_odps_type(
     arrow_type: pa.DataType, col_name: str, unknown_as_string: bool = False
@@ -167,7 +195,51 @@ def odps_schema_to_pandas_dtypes(
     arrow_schema = odps_schema_to_arrow_schema(
         odps_schema, with_partitions=with_partitions
     )
-    return arrow_schema.empty_table().to_pandas().dtypes
+    return arrow_table_to_pandas_dataframe(arrow_schema.empty_table()).dtypes
+
+
+def arrow_table_to_pandas_dataframe(
+    table: pa.Table, meta: DataFrameTableMeta = None
+) -> pd.DataFrame:
+    df = table.to_pandas(
+        types_mapper=lambda x: (
+            ArrowDtype(x) if is_based_for_pandas_dtype(x) else None
+        ),
+        ignore_metadata=True,
+    )
+    if not meta:
+        return df
+
+    # If meta is passed, we should convert the dtypes to match the ones in the meta
+    converted_column_dtypes = dict()
+    for source_dtype, target_col, target_dtype in zip(
+        df.dtypes.values,
+        df.columns,
+        list(meta.pd_index_dtypes.values) + list(meta.pd_column_dtypes.values),
+    ):
+        if source_dtype != target_dtype:
+            # Converting tz-aware dtype to tz-native dtype is a special case.
+            # In numpy1.19, we can't use numpy.dtype.DateTime64Dtype
+            if (
+                isinstance(source_dtype, pd.DatetimeTZDtype)
+                and isinstance(target_dtype, np.dtype)
+                and target_dtype.name.startswith("datetime64")
+            ):
+                df[target_col] = df[target_col].dt.tz_localize(None)
+            else:
+                converted_column_dtypes[target_col] = target_dtype
+
+    if converted_column_dtypes:
+        df = df.astype(converted_column_dtypes)
+
+    return df
+
+
+def pandas_dataframe_to_arrow_table(df: pd.DataFrame, nthreads=1) -> pa.Table:
+    schema = pandas_types_to_arrow_schema(df)
+    return pa.Table.from_pandas(
+        df, schema=schema, nthreads=nthreads, preserve_index=False
+    )
 
 
 def is_scalar_object(df_obj: Any) -> bool:

maxframe/io/odpsio/tableio.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.
@@ -15,6 +15,7 @@
 import os
 import time
 from abc import ABC, abstractmethod
+from collections import OrderedDict
 from contextlib import contextmanager
 from typing import Dict, List, Optional, Union
 
@@ -25,7 +26,7 @@ from odps.apis.storage_api import (
     TableBatchScanResponse,
     TableBatchWriteResponse,
 )
-from odps.tunnel import TableTunnel
+from odps.tunnel import TableDownloadSession, TableDownloadStatus, TableTunnel
 from odps.types import OdpsSchema, PartitionSpec, timestamp_ntz
 from odps.utils import call_with_retry
 
@@ -36,12 +37,13 @@ except ImportError:
 
 from ...config import options
 from ...env import ODPS_STORAGE_API_ENDPOINT
-from ...utils import sync_pyodps_options
+from ...utils import is_empty, sync_pyodps_options
 from .schema import odps_schema_to_arrow_schema
 
 PartitionsType = Union[List[str], str, None]
 
 _DEFAULT_ROW_BATCH_SIZE = 4096
+_DOWNLOAD_ID_CACHE_SIZE = 100
 
 
 class ODPSTableIO(ABC):
@@ -65,7 +67,11 @@ class ODPSTableIO(ABC):
     ) -> OdpsSchema:
         final_cols = []
 
-        columns = columns or [col.name for col in table_schema.simple_columns]
+        columns = (
+            columns
+            if not is_empty(columns)
+            else [col.name for col in table_schema.simple_columns]
+        )
         if partition_columns is True:
             partition_columns = [c.name for c in table_schema.partitions]
         else:
@@ -215,6 +221,46 @@ class TunnelMultiPartitionReader:
 
 
 class TunnelTableIO(ODPSTableIO):
+    _down_session_ids = OrderedDict()
+
+    @classmethod
+    def create_download_sessions(
+        cls,
+        odps_entry: ODPS,
+        full_table_name: str,
+        partitions: List[Optional[str]] = None,
+    ) -> Dict[Optional[str], TableDownloadSession]:
+        table = odps_entry.get_table(full_table_name)
+        tunnel = TableTunnel(odps_entry, quota_name=options.tunnel_quota_name)
+        parts = (
+            [partitions]
+            if partitions is None or isinstance(partitions, str)
+            else partitions
+        )
+        part_to_session = dict()
+        for part in parts:
+            part_key = (full_table_name, part)
+            down_session = None
+
+            if part_key in cls._down_session_ids:
+                down_id = cls._down_session_ids[part_key]
+                down_session = tunnel.create_download_session(
+                    table, async_mode=True, partition_spec=part, download_id=down_id
+                )
+                if down_session.status != TableDownloadStatus.Normal:
+                    down_session = None
+
+            if down_session is None:
+                down_session = tunnel.create_download_session(
+                    table, async_mode=True, partition_spec=part
+                )
+
+            while len(cls._down_session_ids) >= _DOWNLOAD_ID_CACHE_SIZE:
+                cls._down_session_ids.popitem(False)
+            cls._down_session_ids[part_key] = down_session.id
+            part_to_session[part] = down_session
+        return part_to_session
+
     @contextmanager
     def open_reader(
         self,
@@ -241,21 +287,15 @@ class TunnelTableIO(ODPSTableIO):
             or (reverse_range and start is None)
         ):
             with sync_pyodps_options():
-                table = self._odps.get_table(full_table_name)
-                tunnel = TableTunnel(self._odps)
-                parts = (
-                    [partitions]
-                    if partitions is None or isinstance(partitions, str)
-                    else partitions
+                tunnel_sessions = self.create_download_sessions(
+                    self._odps, full_table_name, partitions
+                )
+                part_to_down_id = {
+                    pt: session.id for (pt, session) in tunnel_sessions.items()
+                }
+                total_records = sum(
+                    session.count for session in tunnel_sessions.values()
                 )
-                part_to_down_id = dict()
-                total_records = 0
-                for part in parts:
-                    down_session = tunnel.create_download_session(
-                        table, async_mode=True, partition_spec=part
-                    )
-                    part_to_down_id[part] = down_session.id
-                    total_records += down_session.count
 
         count = None
         if start is not None or stop is not None:
@@ -521,7 +561,10 @@ class HaloTableIO(ODPSTableIO):
 
         table = self._odps.get_table(full_table_name)
         client = StorageApiArrowClient(
-            self._odps, table, rest_endpoint=self._storage_api_endpoint
+            self._odps,
+            table,
+            rest_endpoint=self._storage_api_endpoint,
+            quota_name=options.tunnel_quota_name,
         )
 
         split_option = SplitOptions.SplitMode.SIZE
@@ -597,7 +640,10 @@ class HaloTableIO(ODPSTableIO):
 
         table = self._odps.get_table(full_table_name)
         client = StorageApiArrowClient(
-            self._odps, table, rest_endpoint=self._storage_api_endpoint
+            self._odps,
+            table,
+            rest_endpoint=self._storage_api_endpoint,
+            quota_name=options.tunnel_quota_name,
         )
 
         part_strs = self._convert_partitions(partition)

maxframe/io/odpsio/tests/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/io/odpsio/tests/test_arrow.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.
@@ -14,6 +14,10 @@
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
+import pytest
+
+from maxframe.lib.dtypes_extension import dict_
 
 from ..arrow import arrow_to_pandas, pandas_to_arrow
 
@@ -86,3 +90,44 @@ def test_scalar_convert():
 
     scalar_res = arrow_to_pandas(arrow_data, meta)
     assert scalar_data == scalar_res
+
+
+@pytest.mark.skipif(
+    pa is None or not hasattr(pd, "ArrowDtype"),
+    reason="pandas doesn't support ArrowDtype",
+)
+def test_map_convert():
+    pd_data = pd.DataFrame(
+        {
+            "A": pd.Series(
+                [(("k1", "v1"), ("k2", "v2"))], dtype=dict_(pa.string(), pa.string())
+            ),
+            "B": pd.Series([{"k1": 1, "k2": 2}], dtype=dict_(pa.string(), pa.int64())),
+        },
+    )
+    arrow_data, meta = pandas_to_arrow(pd_data)
+    assert arrow_data.column_names == ["_idx_0", "a", "b"]
+    pd.testing.assert_series_equal(
+        meta.pd_column_dtypes,
+        pd.Series(
+            [dict_(pa.string(), pa.string()), dict_(pa.string(), pa.int64())],
+            index=["A", "B"],
+        ),
+    )
+    pd_result = arrow_to_pandas(arrow_data, meta)
+    pd.testing.assert_frame_equal(pd_data, pd_result)
+
+
+def test_datetime_with_tz_convert():
+    pd_data = pd.DataFrame(
+        {
+            "a": pd.to_datetime(
+                pd.Series([1609459200, 1609545600], index=[0, 1]), unit="s"
+            ),
+        },
+    )
+
+    arrow_data, meta = pandas_to_arrow(pd_data)
+    assert arrow_data.column_names == ["_idx_0", "a"]
+    pd_result = arrow_to_pandas(arrow_data, meta)
+    pd.testing.assert_frame_equal(pd_data, pd_result)

maxframe/io/odpsio/tests/test_schema.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.
@@ -21,6 +21,7 @@ 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 ....utils import pd_release_version
 from ..schema import (
     arrow_schema_to_odps_schema,
@@ -28,6 +29,7 @@ from ..schema import (
     build_table_column_name,
     odps_schema_to_arrow_schema,
     pandas_to_odps_schema,
+    pandas_types_to_arrow_schema,
 )
 
 
@@ -332,3 +334,19 @@ def test_table_meta_with_datetime():
     multiidx = md.Index(raw_multiindex)
     schema, _ = pandas_to_odps_schema(multiidx, unknown_as_string=True)
     assert schema.columns[1].type == odps_types.datetime
+
+
+@pytest.mark.skipif(
+    pa is None or not hasattr(pd, "ArrowDtype"),
+    reason="pandas doesn't support ArrowDtype",
+)
+def test_pandas_types_to_arrow_schema():
+    pd_data = pd.DataFrame(
+        {
+            "int8": pd.Series([], dtype=np.int8),
+            "map": pd.Series([], dtype=dict_(pa.string(), 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())

maxframe/io/odpsio/tests/test_tableio.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/io/odpsio/tests/test_volumeio.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/io/odpsio/volumeio.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/graph/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/graph/connected_components.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/graph/tests/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/graph/tests/test_connected_components.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/llm/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/llm/core.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/llm/models/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/llm/models/dashscope.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/llm/multi_modal.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/llm/text.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/pytorch/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/pytorch/run_function.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/pytorch/run_script.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/pytorch/tests/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/pytorch/tests/test_pytorch.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/utils.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/xgboost/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/xgboost/classifier.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/xgboost/core.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/xgboost/dmatrix.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.
@@ -19,7 +19,7 @@ from ....core.operator.base import Operator
 from ....core.operator.core import TileableOperatorMixin
 from ....dataframe.core import DATAFRAME_TYPE
 from ....serialization.serializables import Float64Field, KeyField, ListField
-from ....serialization.serializables.field import AnyField, Int64Field
+from ....serialization.serializables.field import AnyField, BoolField, Int64Field
 from ....tensor import tensor as astensor
 from ....tensor.core import TENSOR_TYPE
 from ....typing_ import TileableType
@@ -42,6 +42,7 @@ class ToDMatrix(Operator, TileableOperatorMixin):
     qid = AnyField("qid", default=None)
     label_lower_bound = AnyField("label_lower_bound", default=None)
     label_upper_bound = AnyField("label_upper_bound", default=None)
+    enable_categorical = BoolField("enable_categorical", default=None)
 
     @property
     def output_limit(self):
@@ -116,6 +117,7 @@ def to_dmatrix(
     qid=None,
     label_lower_bound=None,
     label_upper_bound=None,
+    enable_categorical=None,
 ):
     data = check_data(data)
     label = check_array_like(label, "label")
@@ -139,6 +141,7 @@ def to_dmatrix(
         label_lower_bound=label_lower_bound,
         label_upper_bound=label_upper_bound,
         gpu=data.op.gpu,
+        enable_categorical=enable_categorical,
         _output_types=get_output_types(data),
     )
     return op()

maxframe/learn/contrib/xgboost/predict.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/xgboost/regressor.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.

maxframe/learn/contrib/xgboost/tests/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
+# 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.