diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py index 629d12311c89..f9661c22f940 100644 --- a/python-package/lightgbm/basic.py +++ b/python-package/lightgbm/basic.py @@ -35,11 +35,14 @@ arrow_is_floating, arrow_is_integer, concat, + pa_array, pa_Array, pa_chunked_array, pa_ChunkedArray, pa_compute, + pa_DataType, pa_Table, + pd_ArrowDtype, pd_CategoricalDtype, pd_DataFrame, pd_Series, @@ -409,6 +412,23 @@ def _is_2d_collection(data: Any) -> bool: return _is_numpy_2d_array(data) or _is_2d_list(data) or isinstance(data, pd_DataFrame) +def _is_pd_arrow_dtype(dtype: Any) -> bool: + """Check whether the dtype is a pandas ArrowDtype.""" + return pd_ArrowDtype is not None and isinstance(dtype, pd_ArrowDtype) + + +def _is_arrow_backed_pd_series(data: Any) -> bool: + """Check whether data is a pandas Series backed by ArrowDtype.""" + return isinstance(data, pd_Series) and _is_pd_arrow_dtype(data.dtype) + + +def _extract_arrow_array(data: Any) -> Any: + """Extract PyArrow array from arrow-backed pandas Series, or return as-is if already an arrow array.""" + if _is_arrow_backed_pd_series(data): + return data.array.__arrow_array__() + return data + + def _is_pyarrow_array(data: Any) -> "TypeGuard[Union[pa_Array, pa_ChunkedArray]]": """Check whether data is a PyArrow array.""" return isinstance(data, (pa_Array, pa_ChunkedArray)) @@ -790,6 +810,55 @@ def _c_int_array(data: np.ndarray) -> Tuple[_ctypes_int_ptr, int, np.ndarray]: return (ptr_data, type_data, data) # return `data` to avoid the temporary copy is freed +def _is_allowed_arrow_dtype(arrow_type: pa_DataType) -> bool: + """Check if the Arrow type is one that C++ can handle (integer, floating point, or boolean).""" + return arrow_is_integer(arrow_type) or arrow_is_floating(arrow_type) or arrow_is_boolean(arrow_type) + + +def _check_for_bad_arrow_dtypes(table: pa_Table) -> None: + """Raise if the Arrow table contains non-numeric types.""" + bad_arrow_dtypes = [ + f"{field.name}: {field.type}" + for field in table.schema + if not _is_allowed_arrow_dtype(field.type) + ] + if bad_arrow_dtypes: + raise ValueError( + "Arrow columns must be int, float or bool.\n" + f"Fields with bad Arrow dtypes: {', '.join(bad_arrow_dtypes)}" + ) + + +def _pandas_to_arrow_table(data: pd_DataFrame) -> pa_Table: + """Convert a pandas DataFrame to an Arrow Table.""" + if not (PYARROW_INSTALLED and CFFI_INSTALLED): + raise LightGBMError("Cannot init Dataset from Arrow without 'pyarrow' and 'cffi' installed.") + arrays = [] + names = [] + for col_name, col in data.items(): + dtype = col.dtype + if _is_pd_arrow_dtype(dtype): + # arrow-backed columns can be converted zero-copy by extracting the underlying array + arrays.append(col.array.__arrow_array__()) + elif hasattr(dtype, 'numpy_dtype'): + # nullable extension types (e.g. pd.Int64Dtype, pd.Float64Dtype, pd.BooleanDtype) + # use separate data+mask arrays; pass the Series directly so pyarrow can interpret the mask + arrays.append(pa_array(col, from_pandas=True)) + elif _is_allowed_numpy_dtype(dtype.type): + # numpy-backed columns are a single contiguous buffer; .values is a zero-copy view + arrays.append(pa_array(col.values, from_pandas=True)) + else: + raise ValueError( + f"Column '{col_name}' has unsupported dtype '{dtype}'. " + "Columns must have integer, floating point, or boolean dtypes." + ) + names.append(str(col_name)) + # column names are coerced to strings as required by Arrow + table = pa_Table.from_arrays(arrays, names=names) + _check_for_bad_arrow_dtypes(table) + return table + + def _is_allowed_numpy_dtype(dtype: type) -> bool: float128 = getattr(np, "float128", type(None)) return issubclass(dtype, (np.integer, np.floating, np.bool_)) and not issubclass(dtype, (np.timedelta64, float128)) @@ -829,7 +898,7 @@ def _data_from_pandas( feature_name: _LGBM_FeatureNameConfiguration, categorical_feature: _LGBM_CategoricalFeatureConfiguration, pandas_categorical: Optional[List[List]], -) -> Tuple[np.ndarray, List[str], Union[List[str], List[int]], List[List]]: +) -> Tuple[Union[np.ndarray, pa_Table], List[str], Union[List[str], List[int]], List[List]]: if len(data.shape) != 2 or data.shape[0] < 1: raise ValueError("Input data must be 2 dimensional and non empty.") @@ -859,6 +928,14 @@ def _data_from_pandas( if categorical_feature == "auto": categorical_feature = cat_cols_not_ordered + if any(_is_pd_arrow_dtype(dtype) for dtype in data.dtypes): + return ( + _pandas_to_arrow_table(data), + feature_name, + categorical_feature, + pandas_categorical, + ) + df_dtypes = [dtype.type for dtype in data.dtypes] # so that the target dtype considers floats df_dtypes.append(np.float32) @@ -1718,9 +1795,7 @@ def __pred_for_pyarrow_table( if not (PYARROW_INSTALLED and CFFI_INSTALLED): raise LightGBMError("Cannot predict from Arrow without 'pyarrow' and 'cffi' installed.") - # Check that the input is valid: we only handle numbers (for now) - if not all(arrow_is_integer(t) or arrow_is_floating(t) or arrow_is_boolean(t) for t in table.schema.types): - raise ValueError("Arrow table may only have integer or floating point datatypes") + _check_for_bad_arrow_dtypes(table) # Prepare prediction output array n_preds = self.__get_num_preds( @@ -2473,9 +2548,7 @@ def __init_from_pyarrow_table( if not (PYARROW_INSTALLED and CFFI_INSTALLED): raise LightGBMError("Cannot init Dataset from Arrow without 'pyarrow' and 'cffi' installed.") - # Check that the input is valid: we only handle numbers (for now) - if not all(arrow_is_integer(t) or arrow_is_floating(t) or arrow_is_boolean(t) for t in table.schema.types): - raise ValueError("Arrow table may only have integer or floating point datatypes") + _check_for_bad_arrow_dtypes(table) # Export Arrow table to C c_array = _export_arrow_to_c(table) @@ -2804,6 +2877,8 @@ def set_field( ) return self + data = _extract_arrow_array(data) + # If the data is a arrow data, we can just pass it to C if _is_pyarrow_array(data) or _is_pyarrow_table(data): # If a table is being passed, we concatenate the columns. This is only valid for @@ -3076,6 +3151,7 @@ def set_label(self, label: Optional[_LGBM_LabelType]) -> "Dataset": ---------- label : list, numpy 1-D array, pandas Series / one-column DataFrame, pyarrow Array, pyarrow ChunkedArray or None The label information to be set into Dataset. + Supports arrow-backed pandas Series (ArrowDtype). Returns ------- @@ -3088,7 +3164,7 @@ def set_label(self, label: Optional[_LGBM_LabelType]) -> "Dataset": if len(label.columns) > 1: raise ValueError("DataFrame for label cannot have multiple columns") label_array = np.ravel(_pandas_to_numpy(label, target_dtype=np.float32)) - elif _is_pyarrow_array(label): + elif _is_pyarrow_array(label) or _is_arrow_backed_pd_series(label): label_array = label else: label_array = _list_to_1d_numpy(data=label, dtype=np.float32, name="label") @@ -3106,6 +3182,7 @@ def set_weight( ---------- weight : list, numpy 1-D array, pandas Series, pyarrow Array, pyarrow ChunkedArray or None Weight to be set for each data point. Weights should be non-negative. + Supports arrow-backed pandas Series (ArrowDtype). Returns ------- @@ -3114,6 +3191,7 @@ def set_weight( """ # Check if the weight contains values other than one if weight is not None: + weight = _extract_arrow_array(weight) if _is_pyarrow_array(weight): # TODO: remove 'type: ignore[attr-defined]' when https://github.com/apache/arrow/issues/49831 is resolved. if pa_compute.all(pa_compute.equal(weight, 1)).as_py(): # type: ignore[attr-defined] @@ -3124,7 +3202,7 @@ def set_weight( # Set field if self._handle is not None and weight is not None: - if not _is_pyarrow_array(weight): + if not (_is_pyarrow_array(weight) or _is_arrow_backed_pd_series(weight)): weight = _list_to_1d_numpy(data=weight, dtype=np.float32, name="weight") self.set_field("weight", weight) self.weight = self.get_field("weight") # original values can be modified at cpp side @@ -3140,6 +3218,7 @@ def set_init_score( ---------- init_score : list, list of lists (for multi-class task), numpy array, pandas Series, pandas DataFrame (for multi-class task), pyarrow Array, pyarrow ChunkedArray, pyarrow Table (for multi-class task) or None Init score for Booster. + Supports arrow-backed pandas Series and DataFrame (ArrowDtype). Returns ------- @@ -3166,6 +3245,7 @@ def set_group( sum(group) = n_samples. For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups, where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc. + Supports arrow-backed pandas Series (ArrowDtype). Returns ------- @@ -3174,7 +3254,7 @@ def set_group( """ self.group = group if self._handle is not None and group is not None: - if not _is_pyarrow_array(group): + if not (_is_pyarrow_array(group) or _is_arrow_backed_pd_series(group)): group = _list_to_1d_numpy(data=group, dtype=np.int32, name="group") self.set_field("group", group) # original values can be modified at cpp side @@ -3191,8 +3271,9 @@ def set_position( Parameters ---------- - position : numpy 1-D array, pandas Series or None, optional (default=None) + position : numpy 1-D array, pandas Series, pyarrow Array, pyarrow ChunkedArray or None, optional (default=None) Position of items used in unbiased learning-to-rank task. + Supports arrow-backed pandas Series (ArrowDtype). Returns ------- @@ -3201,7 +3282,8 @@ def set_position( """ self.position = position if self._handle is not None and position is not None: - position = _list_to_1d_numpy(data=position, dtype=np.int32, name="position") + if not (_is_pyarrow_array(position) or _is_arrow_backed_pd_series(position)): + position = _list_to_1d_numpy(data=position, dtype=np.int32, name="position") self.set_field("position", position) return self diff --git a/python-package/lightgbm/compat.py b/python-package/lightgbm/compat.py index e836b95e9293..4658e11a9500 100644 --- a/python-package/lightgbm/compat.py +++ b/python-package/lightgbm/compat.py @@ -176,6 +176,10 @@ class _LGBMRegressorBase: # type: ignore from pandas import CategoricalDtype as pd_CategoricalDtype except ImportError: from pandas.api.types import CategoricalDtype as pd_CategoricalDtype + try: + from pandas import ArrowDtype as pd_ArrowDtype + except ImportError: + pd_ArrowDtype = None PANDAS_INSTALLED = True except ImportError: PANDAS_INSTALLED = False @@ -198,6 +202,12 @@ class pd_CategoricalDtype: # type: ignore def __init__(self, *args: Any, **kwargs: Any): pass + class pd_ArrowDtype: # type: ignore + """Dummy class for pandas.ArrowDtype.""" + + def __init__(self, *args: Any, **kwargs: Any): + pass + concat = None """matplotlib""" @@ -285,6 +295,7 @@ def __init__(self, *args: Any, **kwargs: Any): from pyarrow import Table as pa_Table from pyarrow import array as pa_array from pyarrow import chunked_array as pa_chunked_array + from pyarrow.lib import DataType as pa_DataType from pyarrow.types import is_boolean as arrow_is_boolean from pyarrow.types import is_floating as arrow_is_floating from pyarrow.types import is_integer as arrow_is_integer @@ -311,6 +322,12 @@ class pa_Table: # type: ignore def __init__(self, *args: Any, **kwargs: Any): pass + class pa_DataType: # type: ignore + """Dummy class for pa.lib.DataType.""" + + def __init__(self, *args: Any, **kwargs: Any): + pass + class pa_compute: # type: ignore """Dummy class for pyarrow.compute module.""" diff --git a/tests/python_package_test/test_basic.py b/tests/python_package_test/test_basic.py index 0d899fde43c7..56032b4dd7c4 100644 --- a/tests/python_package_test/test_basic.py +++ b/tests/python_package_test/test_basic.py @@ -14,10 +14,22 @@ from sklearn.model_selection import train_test_split import lightgbm as lgb -from lightgbm.compat import PANDAS_INSTALLED, pd_DataFrame, pd_Series +from lightgbm.compat import PANDAS_INSTALLED, PYARROW_INSTALLED, pd_DataFrame, pd_Series from .utils import dummy_obj, load_breast_cancer, mse_obj, np_assert_array_equal +# The "name[pyarrow]" dtype-string syntax used by some tests below requires +# pandas >= 1.5 (which introduced ``pd.ArrowDtype``). The "oldest" CI matrix +# pins pandas to a version that does not understand those strings even when +# pyarrow itself is installed, so we gate those tests on this capability. +if PANDAS_INSTALLED and PYARROW_INSTALLED: + import pandas as _pd_for_arrow_check + + _PANDAS_ARROW_DTYPE_SUPPORTED = hasattr(_pd_for_arrow_check, "ArrowDtype") + del _pd_for_arrow_check +else: + _PANDAS_ARROW_DTYPE_SUPPORTED = False + def test_basic(tmp_path): X_train, X_test, y_train, y_test = train_test_split( @@ -1209,3 +1221,202 @@ def test_refit_correctly_handles_categorical_features_in_params(rng) -> None: match=re.escape("Using refit() to change which columns are treated as categorical is not supported"), ): loaded_bst_new = loaded_bst.refit(X_new, y_new, categorical_feature=[0, 1]) + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +@pytest.mark.parametrize( + "dtype_str", + [ + "int8[pyarrow]", + "int16[pyarrow]", + "int32[pyarrow]", + "int64[pyarrow]", + "uint8[pyarrow]", + "uint32[pyarrow]", + "uint64[pyarrow]", + "float32[pyarrow]", + "float64[pyarrow]", + "bool[pyarrow]", + ] +) +def test_dataset_constructor_with_arrow_backed_dataframe(dtype_str, rng): + pd = pytest.importorskip("pandas") + X = pd.DataFrame(rng.integers(0, 10, (100, 5)), dtype=dtype_str) + ds = lgb.Dataset(X).construct() + assert ds.num_data() == 100 + assert ds.num_feature() == 5 + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_dataset_constructor_with_arrow_backed_dataframe_mixed(rng): + """tests for mixes of both numpy series and pyarrow-backed sereis in the same dataframe""" + pd = pytest.importorskip("pandas") + X = pd.DataFrame(rng.integers(0, 10, (100, 5)), dtype="int32[pyarrow]") + X["int_col"] = pd.Series(rng.integers(0, 10, 100), dtype=np.int32) + X["float_col"] = pd.Series(rng.uniform(size=100), dtype=np.float32) + X["bool_col"] = pd.Series(rng.choice(a=[False, True], size=100), dtype=np.bool_) + X["nullable_int_col"] = pd.Series(rng.choice(a=[0, 1, None], size=100), dtype=pd.Int32Dtype()) + X["nullable_float_col"] = pd.Series(rng.choice(a=[0.0, 1.0, None], size=100), dtype=pd.Float32Dtype()) + X["nullable_bool_col"] = pd.Series(rng.choice(a=[False, True, None], size=100), dtype=pd.BooleanDtype()) + ds = lgb.Dataset(X).construct() + assert ds.num_data() == 100 + assert ds.num_feature() == 11 + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_dataset_constructor_with_invalid_arrow_backed_dataframe(): + pd = pytest.importorskip("pandas") + for dtype in ["string[pyarrow]", "large_string[pyarrow]", "binary[pyarrow]", "large_binary[pyarrow]"]: + with pytest.raises(ValueError, match="must be int, float or bool."): + lgb.Dataset(pd.DataFrame(["hello", "world", None], dtype=dtype)).construct() + for dtype in [ + "timestamp[ns][pyarrow]", + "timestamp[us][pyarrow]", + "timestamp[ms][pyarrow]", + "timestamp[s][pyarrow]", + "timestamp[ns, tz=UTC][pyarrow]", + "date32[pyarrow]", + "date64[pyarrow]", + ]: + with pytest.raises(ValueError, match="must be int, float or bool."): + lgb.Dataset(pd.DataFrame(["2020-01-01", None], dtype=dtype)).construct() + for dtype in ["time32[s][pyarrow]", "time32[ms][pyarrow]", "time64[ns][pyarrow]", "time64[us][pyarrow]"]: + with pytest.raises(ValueError, match="must be int, float or bool."): + lgb.Dataset(pd.DataFrame(["12:00:00", None], dtype=dtype)).construct() + for dtype in ["duration[s][pyarrow]", "duration[ms][pyarrow]", "duration[us][pyarrow]", "duration[ns][pyarrow]"]: + with pytest.raises(ValueError, match="must be int, float or bool."): + lgb.Dataset(pd.DataFrame([1000, None], dtype=dtype)).construct() + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +@pytest.mark.parametrize( + ("field_name", "dtype", "data_fn"), + [ + ("label", "float64[pyarrow]", lambda rng: rng.random(100)), + ("weight", "float64[pyarrow]", lambda rng: rng.random(100)), + ("group", "int32[pyarrow]", lambda _: [20, 30, 50]), + # Position support will be added in https://github.com/lightgbm-org/LightGBM/pull/7260 + # ("position", "int8[pyarrow]", lambda rng: rng.integers(0, 10, size=100)), + ("init_score", "float64[pyarrow]", lambda _: list(range(100))), + ], +) +def test_dataset_constructor_with_arrow_series_metadata_field(rng, field_name, dtype, data_fn): + pd = pytest.importorskip("pandas") + X = rng.random((100, 5)) + data = data_fn(rng) + series = pd.Series(data, dtype=dtype) + ds = lgb.Dataset(X, **{field_name: series}).construct() + retrieved = getattr(ds, f"get_{field_name}")() + assert isinstance(retrieved, np.ndarray) + expected = np.asarray(data) + np.testing.assert_allclose(retrieved, expected) + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_arrow_backed_categorical_features(rng): + pd = pytest.importorskip("pandas") + X = pd.DataFrame({ + 'cat_col': pd.Series(rng.choice([1, 2, 3], size=100), dtype="float64[pyarrow]"), + 'num_col': rng.random(100) + }) + y = rng.integers(0, 2, 100) + ds = lgb.Dataset(X, label=y, categorical_feature=['cat_col']).construct() + assert ds.num_data() == 100 + assert ds.num_feature() == 2 + assert ds.categorical_feature == ['cat_col'] + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_pd_categorical_features_with_arrow_backed_dataframe(rng): + pd = pytest.importorskip("pandas") + X = pd.DataFrame({ + 'cat_col': pd.Categorical(['a', 'b', 'c'] * 30), + 'num_col': rng.random(90) + }) + y = rng.integers(0, 2, 90) + X['num_col'] = X['num_col'].astype('float64[pyarrow]') + ds = lgb.Dataset(X, label=y, categorical_feature=['cat_col']).construct() + assert ds.num_data() == 90 + assert ds.num_feature() == 2 + assert ds.categorical_feature == ['cat_col'] + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_empty_arrow_backed_dataframe_raises(): + pd = pytest.importorskip("pandas") + X = pd.DataFrame(dtype="float64[pyarrow]") + with pytest.raises(ValueError, match="must be 2 dimensional and non empty"): + lgb.Dataset(X).construct() + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_numeric_column_names_with_arrow_backend(rng): + pd = pytest.importorskip("pandas") + X = pd.DataFrame(rng.random((50, 3)), columns=[0, 1, 2], dtype="float64[pyarrow]") + ds = lgb.Dataset(X).construct() + assert ds.feature_name == ['0', '1', '2'] + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_all_null_arrow_column(): + pd = pytest.importorskip("pandas") + X = pd.DataFrame({ + 'col1': pd.Series([1.0, 2.0, 3.0], dtype='float64[pyarrow]'), + 'col2': pd.Series([None, None, None], dtype='float64[pyarrow]') + }) + y = [0, 1, 0] + ds = lgb.Dataset(X, label=y).construct() + assert ds.num_feature() == 2 + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_all_arrow_vs_mixed_dtypes_equivalence(rng): + pd = pytest.importorskip("pandas") + data = rng.random((100, 5)) + y = rng.integers(0, 2, 100) + + X_all_arrow = pd.DataFrame(data, dtype="float64[pyarrow]") + X_mixed = pd.DataFrame(data) + X_mixed[0] = X_mixed[0].astype('float64[pyarrow]') + + ds1 = lgb.Dataset(X_all_arrow, label=y).construct() # fast path + ds2 = lgb.Dataset(X_mixed, label=y).construct() # mixed path + + params = {'seed': 42, 'verbose': -1} + bst1 = lgb.train(params, ds1, num_boost_round=5) + bst2 = lgb.train(params, ds2, num_boost_round=5) + + pred1 = bst1.predict(data) + pred2 = bst2.predict(data) + + np.testing.assert_allclose(pred1, pred2) diff --git a/tests/python_package_test/test_sklearn.py b/tests/python_package_test/test_sklearn.py index 2834a779a49e..b7b4a22dca71 100644 --- a/tests/python_package_test/test_sklearn.py +++ b/tests/python_package_test/test_sklearn.py @@ -17,7 +17,7 @@ from sklearn.datasets import load_svmlight_file, make_blobs, make_multilabel_classification from sklearn.ensemble import StackingClassifier, StackingRegressor from sklearn.metrics import accuracy_score, log_loss, mean_squared_error, r2_score -from sklearn.model_selection import GridSearchCV, RandomizedSearchCV, train_test_split +from sklearn.model_selection import GridSearchCV, RandomizedSearchCV, cross_val_score, train_test_split from sklearn.multioutput import ClassifierChain, MultiOutputClassifier, MultiOutputRegressor, RegressorChain from sklearn.utils.estimator_checks import parametrize_with_checks as sklearn_parametrize_with_checks from sklearn.utils.validation import check_is_fitted @@ -54,6 +54,18 @@ SKLEARN_VERSION_GTE_1_6 = (int(SKLEARN_MAJOR), int(SKLEARN_MINOR)) >= (1, 6) SKLEARN_VERSION_GTE_1_7 = (int(SKLEARN_MAJOR), int(SKLEARN_MINOR)) >= (1, 7) +# The "name[pyarrow]" dtype-string syntax used by some tests below requires +# pandas >= 1.5 (which introduced ``pd.ArrowDtype``). The "oldest" CI matrix +# pins pandas to a version that does not understand those strings even when +# pyarrow itself is installed, so we gate those tests on this capability. +if PANDAS_INSTALLED and PYARROW_INSTALLED: + import pandas as _pd_for_arrow_check + + _PANDAS_ARROW_DTYPE_SUPPORTED = hasattr(_pd_for_arrow_check, "ArrowDtype") + del _pd_for_arrow_check +else: + _PANDAS_ARROW_DTYPE_SUPPORTED = False + decreasing_generator = itertools.count(0, -1) estimator_classes = (lgb.LGBMModel, lgb.LGBMClassifier, lgb.LGBMRegressor, lgb.LGBMRanker) task_to_model_factory = { @@ -2260,3 +2272,53 @@ def test_eval_X_eval_y_eval_set_equivalence(): assert gbm2.evals_result_["valid_0"]["l2"] != gbm2.evals_result_["valid_1"]["l2"], ( "Evaluation results for the 2 validation sets are not different. This might mean they weren't both used." ) + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_lgbm_classifier_with_arrow_backed_dataframe(): + pd = pytest.importorskip("pandas") + X, y = load_breast_cancer(return_X_y=True) + X_arrow = pd.DataFrame(X, dtype="float64[pyarrow]") + y_arrow = pd.Series(y, dtype="int32[pyarrow]") + X_train, X_test, y_train, y_test = train_test_split(X_arrow, y_arrow, test_size=0.2, random_state=42) + clf = lgb.LGBMClassifier(n_estimators=10, verbose=-1, random_state=42) + clf.fit(X_train, y_train) + y_pred = clf.predict(X_test) + y_pred_proba = clf.predict_proba(X_test) + assert y_pred.shape == (len(X_test),) + assert y_pred_proba.shape == (len(X_test), 2) + assert accuracy_score(y_test, y_pred) > 0.9 + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_lgbm_regressor_with_arrow_backed_dataframe(): + pd = pytest.importorskip("pandas") + X, y = make_synthetic_regression() + X_arrow = pd.DataFrame(X, dtype="float32[pyarrow]") + y_arrow = pd.Series(y, dtype="float32[pyarrow]") + X_train, X_test, y_train, _ = train_test_split(X_arrow, y_arrow, test_size=0.2, random_state=42) + reg = lgb.LGBMRegressor(n_estimators=10, verbose=-1, random_state=42) + reg.fit(X_train, y_train) + y_pred = reg.predict(X_test) + assert y_pred.shape == (len(X_test),) + assert np.isfinite(y_pred).all() + + +@pytest.mark.skipif( + not _PANDAS_ARROW_DTYPE_SUPPORTED, + reason="pandas + pyarrow with ArrowDtype string support (pandas >= 1.5) not installed", +) +def test_cross_validation_with_arrow_backed_dataframe(rng): + pd = pytest.importorskip("pandas") + X_arrow = pd.DataFrame(rng.random((100, 5)), dtype="float32[pyarrow]") + y_arrow = pd.Series(rng.integers(0, 2, 100), dtype="int32[pyarrow]") + clf = lgb.LGBMClassifier(n_estimators=5, verbose=-1) + scores = cross_val_score(clf, X_arrow, y_arrow, cv=3) + assert len(scores) == 3 + assert all(0 <= s <= 1 for s in scores)