diff --git a/python-package/xgboost/compat.py b/python-package/xgboost/compat.py index 582b905f3..0c7b88965 100644 --- a/python-package/xgboost/compat.py +++ b/python-package/xgboost/compat.py @@ -81,13 +81,14 @@ else: # pandas try: from pandas import DataFrame, Series - from pandas import MultiIndex + from pandas import MultiIndex, Int64Index from pandas import concat as pandas_concat PANDAS_INSTALLED = True except ImportError: MultiIndex = object + Int64Index = object DataFrame = object Series = object pandas_concat = None diff --git a/python-package/xgboost/core.py b/python-package/xgboost/core.py index 5d4696886..bd3e601c3 100644 --- a/python-package/xgboost/core.py +++ b/python-package/xgboost/core.py @@ -18,7 +18,7 @@ import numpy as np import scipy.sparse from .compat import ( - STRING_TYPES, DataFrame, MultiIndex, py_str, + STRING_TYPES, DataFrame, MultiIndex, Int64Index, py_str, PANDAS_INSTALLED, DataTable, CUDF_INSTALLED, CUDF_DataFrame, CUDF_Series, CUDF_MultiIndex, os_fspath, os_PathLike) @@ -296,6 +296,8 @@ def _maybe_pandas_data(data, feature_names, feature_types): ' '.join([str(x) for x in i]) for i in data.columns ] + elif isinstance(data.columns, Int64Index): + feature_names = list(map(str, data.columns)) else: feature_names = data.columns.format() diff --git a/tests/python/test_with_pandas.py b/tests/python/test_with_pandas.py index cbdce64b1..e91627713 100644 --- a/tests/python/test_with_pandas.py +++ b/tests/python/test_with_pandas.py @@ -83,6 +83,13 @@ class TestPandas(unittest.TestCase): assert dm.num_row() == 3 assert dm.num_col() == 2 + df_int = pd.DataFrame([[1, 1.1], [2, 2.2]], columns=[9, 10]) + dm_int = xgb.DMatrix(df_int) + df_range = pd.DataFrame([[1, 1.1], [2, 2.2]], columns=range(9, 11, 1)) + dm_range = xgb.DMatrix(df_range) + assert dm_int.feature_names == ['9', '10'] # assert not "9 " + assert dm_int.feature_names == dm_range.feature_names + # test MultiIndex as columns df = pd.DataFrame( [