chore: import upstream snapshot with attribution

tests/python_package_test/conftest.py

@@ -0,0 +1,12 @@
+import numpy as np
+import pytest
+
+
+@pytest.fixture(scope="function")
+def rng():
+    return np.random.default_rng()
+
+
+@pytest.fixture(scope="function")
+def rng_fixed_seed():
+    return np.random.default_rng(seed=42)

tests/python_package_test/test_arrow.py

@@ -0,0 +1,512 @@
+# coding: utf-8
+import filecmp
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import numpy as np
+import pytest
+
+import lightgbm as lgb
+
+from .utils import np_assert_array_equal
+
+pa = pytest.importorskip("pyarrow")
+
+
+# ----------------------------------------------------------------------------------------------- #
+#                                            UTILITIES                                            #
+# ----------------------------------------------------------------------------------------------- #
+
+_INTEGER_TYPES = [
+    pa.int8(),
+    pa.int16(),
+    pa.int32(),
+    pa.int64(),
+    pa.uint8(),
+    pa.uint16(),
+    pa.uint32(),
+    pa.uint64(),
+]
+_FLOAT_TYPES = [
+    pa.float32(),
+    pa.float64(),
+]
+
+
+def generate_simple_arrow_table(empty_chunks: bool = False) -> pa.Table:
+    c: list[list[int]] = [[]] if empty_chunks else []
+    columns = [
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint8()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int8()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint16()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int16()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint32()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int32()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint64()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int64()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.float32()),
+        pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.float64()),
+        pa.chunked_array(c + [[True, True, False]] + c + [[False, True]] + c, type=pa.bool_()),
+    ]
+    return pa.Table.from_arrays(columns, names=[f"col_{i}" for i in range(len(columns))])
+
+
+def generate_nullable_arrow_table(dtype: Any) -> pa.Table:
+    columns = [
+        pa.chunked_array([[1, None, 3, 4, 5]], type=dtype),
+        pa.chunked_array([[None, 2, 3, 4, 5]], type=dtype),
+        pa.chunked_array([[1, 2, 3, 4, None]], type=dtype),
+        pa.chunked_array([[None, None, None, None, None]], type=dtype),
+    ]
+    return pa.Table.from_arrays(columns, names=[f"col_{i}" for i in range(len(columns))])
+
+
+def generate_dummy_arrow_table() -> pa.Table:
+    col1 = pa.chunked_array([[1, 2, 3], [4, 5]], type=pa.uint8())
+    col2 = pa.chunked_array([[0.5, 0.6], [0.1, 0.8, 1.5]], type=pa.float32())
+    return pa.Table.from_arrays([col1, col2], names=["a", "b"])
+
+
+def generate_random_arrow_table(
+    num_columns: int,
+    num_datapoints: int,
+    seed: int,
+    generate_nulls: bool = True,
+    values: Optional[np.ndarray] = None,
+) -> pa.Table:
+    columns = [
+        generate_random_arrow_array(num_datapoints, seed + i, generate_nulls=generate_nulls, values=values)
+        for i in range(num_columns)
+    ]
+    names = [f"col_{i}" for i in range(num_columns)]
+    return pa.Table.from_arrays(columns, names=names)
+
+
+def generate_random_arrow_array(
+    num_datapoints: int,
+    seed: int,
+    generate_nulls: bool = True,
+    values: Optional[np.ndarray] = None,
+) -> pa.ChunkedArray:
+    generator = np.random.default_rng(seed)
+    data = (
+        generator.standard_normal(num_datapoints)
+        if values is None
+        else generator.choice(values, size=num_datapoints, replace=True)
+    )
+
+    # Set random nulls
+    if generate_nulls:
+        indices = generator.choice(len(data), size=num_datapoints // 10)
+        data[indices] = None
+
+    # Split data into <=2 random chunks
+    split_points = np.sort(generator.choice(np.arange(1, num_datapoints), 2, replace=False))
+    split_points = np.concatenate([[0], split_points, [num_datapoints]])
+    chunks = [data[split_points[i] : split_points[i + 1]] for i in range(len(split_points) - 1)]
+    chunks = [chunk for chunk in chunks if len(chunk) > 0]
+
+    # Turn chunks into array
+    return pa.chunked_array(chunks, type=pa.float32())
+
+
+def dummy_dataset_params() -> Dict[str, Any]:
+    return {
+        "min_data_in_bin": 1,
+        "min_data_in_leaf": 1,
+    }
+
+
+# ----------------------------------------------------------------------------------------------- #
+#                                            UNIT TESTS                                           #
+# ----------------------------------------------------------------------------------------------- #
+
+# ------------------------------------------- DATASET ------------------------------------------- #
+
+
+def assert_datasets_equal(tmp_path: Path, lhs: lgb.Dataset, rhs: lgb.Dataset):
+    lhs._dump_text(tmp_path / "arrow.txt")
+    rhs._dump_text(tmp_path / "pandas.txt")
+    assert filecmp.cmp(tmp_path / "arrow.txt", tmp_path / "pandas.txt")
+
+
+@pytest.mark.parametrize(
+    ("arrow_table_fn", "dataset_params"),
+    [  # Use lambda functions here to minimize memory consumption
+        (generate_simple_arrow_table, dummy_dataset_params()),
+        (lambda: generate_simple_arrow_table(empty_chunks=True), dummy_dataset_params()),
+        (generate_dummy_arrow_table, dummy_dataset_params()),
+        (lambda: generate_nullable_arrow_table(pa.float32()), dummy_dataset_params()),
+        (lambda: generate_nullable_arrow_table(pa.int32()), dummy_dataset_params()),
+        (lambda: generate_random_arrow_table(3, 1000, 42), {}),
+        (lambda: generate_random_arrow_table(100, 10000, 43), {}),
+    ],
+)
+def test_dataset_construct_fuzzy(tmp_path, arrow_table_fn, dataset_params):
+    arrow_table = arrow_table_fn()
+
+    arrow_dataset = lgb.Dataset(arrow_table, params=dataset_params)
+    arrow_dataset.construct()
+
+    pandas_dataset = lgb.Dataset(arrow_table.to_pandas(), params=dataset_params)
+    pandas_dataset.construct()
+
+    assert_datasets_equal(tmp_path, arrow_dataset, pandas_dataset)
+
+
+def test_dataset_construct_fuzzy_boolean(tmp_path):
+    boolean_data = generate_random_arrow_table(10, 10000, 42, generate_nulls=False, values=np.array([True, False]))
+
+    float_schema = pa.schema([pa.field(f"col_{i}", pa.float32()) for i in range(len(boolean_data.columns))])
+    float_data = boolean_data.cast(float_schema)
+
+    arrow_dataset = lgb.Dataset(boolean_data)
+    arrow_dataset.construct()
+
+    pandas_dataset = lgb.Dataset(float_data.to_pandas())
+    pandas_dataset.construct()
+
+    assert_datasets_equal(tmp_path, arrow_dataset, pandas_dataset)
+
+
+# -------------------------------------------- FIELDS ------------------------------------------- #
+
+
+def test_dataset_construct_fields_fuzzy():
+    arrow_table = generate_random_arrow_table(3, 1000, 42)
+    arrow_labels = generate_random_arrow_array(1000, 42, generate_nulls=False)
+    arrow_weights = generate_random_arrow_array(1000, 42, generate_nulls=False)
+    arrow_groups = pa.chunked_array([[300, 400, 50], [250]], type=pa.int32())
+
+    arrow_dataset = lgb.Dataset(arrow_table, label=arrow_labels, weight=arrow_weights, group=arrow_groups)
+    arrow_dataset.construct()
+
+    pandas_dataset = lgb.Dataset(
+        arrow_table.to_pandas(),
+        label=arrow_labels.to_numpy(),
+        weight=arrow_weights.to_numpy(),
+        group=arrow_groups.to_numpy(),
+    )
+    pandas_dataset.construct()
+
+    # Check for equality
+    for field in ("label", "weight", "group"):
+        np_assert_array_equal(arrow_dataset.get_field(field), pandas_dataset.get_field(field), strict=True)
+    np_assert_array_equal(arrow_dataset.get_label(), pandas_dataset.get_label(), strict=True)
+    np_assert_array_equal(arrow_dataset.get_weight(), pandas_dataset.get_weight(), strict=True)
+
+
+# -------------------------------------------- LABELS ------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "label_data",
+    [
+        [[0, 1, 0, 0, 1]],
+        [[0], [1, 0, 0, 1]],
+        [[], [0], [1, 0, 0, 1]],
+        [[0], [], [1, 0], [], [], [0, 1], []],
+    ],
+)
+@pytest.mark.parametrize("arrow_type", _INTEGER_TYPES + _FLOAT_TYPES)
+def test_dataset_construct_labels(label_data, arrow_type):
+    data = generate_dummy_arrow_table()
+    labels = pa.chunked_array(label_data, type=arrow_type)
+    dataset = lgb.Dataset(data, label=labels, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 0, 0, 1], dtype=np.float32)
+    np_assert_array_equal(expected, dataset.get_label(), strict=True)
+
+
+@pytest.mark.parametrize(
+    "label_data",
+    [
+        [[False, True, False, False, True]],
+        [[False], [True, False, False, True]],
+        [[], [False], [True, False, False, True]],
+        [[False], [], [True, False], [], [], [False, True], []],
+    ],
+)
+def test_dataset_construct_labels_boolean(label_data):
+    data = generate_dummy_arrow_table()
+    labels = pa.chunked_array(label_data, type=pa.bool_())
+    dataset = lgb.Dataset(data, label=labels, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 0, 0, 1], dtype=np.float32)
+    np_assert_array_equal(expected, dataset.get_label(), strict=True)
+
+
+# ------------------------------------------- WEIGHTS ------------------------------------------- #
+
+
+def test_dataset_construct_weights_none():
+    data = generate_dummy_arrow_table()
+    weight = pa.chunked_array([[1, 1, 1, 1, 1]])
+    dataset = lgb.Dataset(data, weight=weight, params=dummy_dataset_params())
+    dataset.construct()
+    assert dataset.get_weight() is None
+    assert dataset.get_field("weight") is None
+
+
+@pytest.mark.parametrize(
+    "weight_data",
+    [
+        [[3, 0.7, 1.5, 0.5, 0.1]],
+        [[3], [0.7, 1.5, 0.5, 0.1]],
+        [[], [3], [0.7, 1.5, 0.5, 0.1]],
+        [[3], [0.7], [], [], [1.5, 0.5, 0.1], []],
+    ],
+)
+@pytest.mark.parametrize("arrow_type", _FLOAT_TYPES)
+def test_dataset_construct_weights(weight_data, arrow_type):
+    data = generate_dummy_arrow_table()
+    weights = pa.chunked_array(weight_data, type=arrow_type)
+    dataset = lgb.Dataset(data, weight=weights, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([3, 0.7, 1.5, 0.5, 0.1], dtype=np.float32)
+    np_assert_array_equal(expected, dataset.get_weight(), strict=True)
+
+
+# -------------------------------------------- GROUPS ------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "group_data",
+    [
+        [[2, 3]],
+        [[2], [3]],
+        [[], [2, 3]],
+        [[2], [], [3], []],
+    ],
+)
+@pytest.mark.parametrize("arrow_type", _INTEGER_TYPES)
+def test_dataset_construct_groups(group_data, arrow_type):
+    data = generate_dummy_arrow_table()
+    groups = pa.chunked_array(group_data, type=arrow_type)
+    dataset = lgb.Dataset(data, group=groups, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 2, 5], dtype=np.int32)
+    np_assert_array_equal(expected, dataset.get_field("group"), strict=True)
+
+
+# ----------------------------------------- INIT SCORES ----------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "init_score_data",
+    [
+        [[0, 1, 2, 3, 3]],
+        [[0, 1, 2], [3, 3]],
+        [[], [0, 1, 2], [3, 3]],
+        [[0, 1], [], [], [2], [3, 3], []],
+    ],
+)
+@pytest.mark.parametrize("arrow_type", _INTEGER_TYPES + _FLOAT_TYPES)
+def test_dataset_construct_init_scores_array(init_score_data, arrow_type):
+    data = generate_dummy_arrow_table()
+    init_scores = pa.chunked_array(init_score_data, type=arrow_type)
+    dataset = lgb.Dataset(data, init_score=init_scores, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 2, 3, 3], dtype=np.float64)
+    np_assert_array_equal(expected, dataset.get_init_score(), strict=True)
+
+
+def test_dataset_construct_init_scores_table():
+    data = generate_dummy_arrow_table()
+    init_scores = pa.Table.from_arrays(
+        [
+            generate_random_arrow_array(5, seed=1, generate_nulls=False),
+            generate_random_arrow_array(5, seed=2, generate_nulls=False),
+            generate_random_arrow_array(5, seed=3, generate_nulls=False),
+        ],
+        names=["a", "b", "c"],
+    )
+    dataset = lgb.Dataset(data, init_score=init_scores, params=dummy_dataset_params())
+    dataset.construct()
+
+    actual = dataset.get_init_score()
+    expected = init_scores.to_pandas().to_numpy().astype(np.float64)
+    np_assert_array_equal(expected, actual, strict=True)
+
+
+# ------------------------------------------ PREDICTION ----------------------------------------- #
+
+
+def assert_equal_predict_arrow_pandas(booster: lgb.Booster, data: pa.Table):
+    p_arrow = booster.predict(data)
+    p_pandas = booster.predict(data.to_pandas())
+    np_assert_array_equal(p_arrow, p_pandas, strict=True)
+
+    p_raw_arrow = booster.predict(data, raw_score=True)
+    p_raw_pandas = booster.predict(data.to_pandas(), raw_score=True)
+    np_assert_array_equal(p_raw_arrow, p_raw_pandas, strict=True)
+
+    p_leaf_arrow = booster.predict(data, pred_leaf=True)
+    p_leaf_pandas = booster.predict(data.to_pandas(), pred_leaf=True)
+    np_assert_array_equal(p_leaf_arrow, p_leaf_pandas, strict=True)
+
+    p_pred_contrib_arrow = booster.predict(data, pred_contrib=True)
+    p_pred_contrib_pandas = booster.predict(data.to_pandas(), pred_contrib=True)
+    np_assert_array_equal(p_pred_contrib_arrow, p_pred_contrib_pandas, strict=True)
+
+    p_first_iter_arrow = booster.predict(data, start_iteration=0, num_iteration=1, raw_score=True)
+    p_first_iter_pandas = booster.predict(data.to_pandas(), start_iteration=0, num_iteration=1, raw_score=True)
+    np_assert_array_equal(p_first_iter_arrow, p_first_iter_pandas, strict=True)
+
+
+def test_predict_regression():
+    data_float = generate_random_arrow_table(10, 10000, 42)
+    data_bool = generate_random_arrow_table(1, 10000, 42, generate_nulls=False, values=np.array([True, False]))
+    data = pa.Table.from_arrays(data_float.columns + data_bool.columns, names=data_float.schema.names + ["col_bool"])
+
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_arrow_array(10000, 43, generate_nulls=False),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "regression", "num_leaves": 7},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_arrow_pandas(booster, data)
+
+
+def test_predict_binary_classification():
+    data = generate_random_arrow_table(10, 10000, 42)
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_arrow_array(10000, 43, generate_nulls=False, values=np.arange(2)),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "binary", "num_leaves": 7},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_arrow_pandas(booster, data)
+
+
+def test_predict_multiclass_classification():
+    data = generate_random_arrow_table(10, 10000, 42)
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_arrow_array(10000, 43, generate_nulls=False, values=np.arange(5)),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "multiclass", "num_leaves": 7, "num_class": 5},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_arrow_pandas(booster, data)
+
+
+def test_predict_ranking():
+    data = generate_random_arrow_table(10, 10000, 42)
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_arrow_array(10000, 43, generate_nulls=False, values=np.arange(4)),
+        group=np.array([1000, 2000, 3000, 4000]),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "lambdarank", "num_leaves": 7},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_arrow_pandas(booster, data)
+
+
+def test_arrow_feature_name_auto():
+    data = generate_dummy_arrow_table()
+    dataset = lgb.Dataset(
+        data,
+        label=pa.chunked_array([[0, 1, 0, 0, 1]]),
+        params=dummy_dataset_params(),
+        categorical_feature=["a"],
+    )
+    booster = lgb.train({"num_leaves": 7}, dataset, num_boost_round=5)
+    assert booster.feature_name() == ["a", "b"]
+
+
+def test_arrow_feature_name_manual():
+    data = generate_dummy_arrow_table()
+    dataset = lgb.Dataset(
+        data,
+        label=pa.chunked_array([[0, 1, 0, 0, 1]]),
+        params=dummy_dataset_params(),
+        feature_name=["c", "d"],
+        categorical_feature=["c"],
+    )
+    booster = lgb.train({"num_leaves": 7}, dataset, num_boost_round=5)
+    assert booster.feature_name() == ["c", "d"]
+
+
+def pyarrow_array_equal(arr1: pa.ChunkedArray, arr2: pa.ChunkedArray) -> bool:
+    """Similar to ``np.array_equal()``, but for ``pyarrow.Array`` objects.
+
+    ``pyarrow.Array`` objects with identical values do not compare equal if any of those
+    values are nulls. This function treats them as equal.
+    """
+    if len(arr1) != len(arr2):
+        return False
+
+    np1 = arr1.to_numpy()
+    np2 = arr2.to_numpy()
+    return np.array_equal(np1, np2, equal_nan=True)
+
+
+def test_get_data_arrow_table():
+    original_table = generate_simple_arrow_table()
+    dataset = lgb.Dataset(original_table, free_raw_data=False)
+    dataset.construct()
+
+    returned_data = dataset.get_data()
+    assert isinstance(returned_data, pa.Table)
+    assert returned_data.schema == original_table.schema
+    assert returned_data.shape == original_table.shape
+
+    for column_name in original_table.column_names:
+        original_column = original_table[column_name]
+        returned_column = returned_data[column_name]
+
+        assert original_column.type == returned_column.type
+        assert original_column.num_chunks == returned_column.num_chunks
+        assert pyarrow_array_equal(original_column, returned_column)
+
+        for i in range(original_column.num_chunks):
+            original_chunk_array = pa.chunked_array([original_column.chunk(i)])
+            returned_chunk_array = pa.chunked_array([returned_column.chunk(i)])
+            assert pyarrow_array_equal(original_chunk_array, returned_chunk_array)
+
+
+def test_get_data_arrow_table_subset(rng):
+    original_table = generate_random_arrow_table(num_columns=3, num_datapoints=1000, seed=42)
+    dataset = lgb.Dataset(original_table, free_raw_data=False)
+    dataset.construct()
+
+    subset_size = 100
+    used_indices = rng.choice(a=original_table.shape[0], size=subset_size, replace=False)
+    used_indices = sorted(used_indices)
+
+    subset_dataset = dataset.subset(used_indices).construct()
+    expected_subset = original_table.take(used_indices)
+    subset_data = subset_dataset.get_data()
+
+    assert isinstance(subset_data, pa.Table)
+    assert subset_data.schema == expected_subset.schema
+    assert subset_data.shape == expected_subset.shape
+    assert len(subset_data) == len(used_indices)
+    assert subset_data.shape == (subset_size, 3)
+
+    for column_name in expected_subset.column_names:
+        expected_col = expected_subset[column_name]
+        returned_col = subset_data[column_name]
+        assert expected_col.type == returned_col.type
+        assert pyarrow_array_equal(expected_col, returned_col)

tests/python_package_test/test_callback.py

@@ -0,0 +1,65 @@
+# coding: utf-8
+import pytest
+
+import lightgbm as lgb
+
+from .utils import SERIALIZERS, pickle_and_unpickle_object
+
+
+def reset_feature_fraction(boosting_round):
+    return 0.6 if boosting_round < 15 else 0.8
+
+
+@pytest.mark.parametrize("serializer", SERIALIZERS)
+def test_early_stopping_callback_is_picklable(serializer):
+    rounds = 5
+    callback = lgb.early_stopping(stopping_rounds=rounds)
+    callback_from_disk = pickle_and_unpickle_object(obj=callback, serializer=serializer)
+    assert callback_from_disk.order == 30
+    assert callback_from_disk.before_iteration is False
+    assert callback.stopping_rounds == callback_from_disk.stopping_rounds
+    assert callback.stopping_rounds == rounds
+
+
+def test_early_stopping_callback_rejects_invalid_stopping_rounds_with_informative_errors():
+    with pytest.raises(TypeError, match="early_stopping_round should be an integer. Got 'str'"):
+        lgb.early_stopping(stopping_rounds="neverrrr")
+
+
+@pytest.mark.parametrize("stopping_rounds", [-10, -1, 0])
+def test_early_stopping_callback_accepts_non_positive_stopping_rounds(stopping_rounds):
+    cb = lgb.early_stopping(stopping_rounds=stopping_rounds)
+    assert cb.enabled is False
+
+
+@pytest.mark.parametrize("serializer", SERIALIZERS)
+def test_log_evaluation_callback_is_picklable(serializer):
+    periods = 42
+    callback = lgb.log_evaluation(period=periods)
+    callback_from_disk = pickle_and_unpickle_object(obj=callback, serializer=serializer)
+    assert callback_from_disk.order == 10
+    assert callback_from_disk.before_iteration is False
+    assert callback.period == callback_from_disk.period
+    assert callback.period == periods
+
+
+@pytest.mark.parametrize("serializer", SERIALIZERS)
+def test_record_evaluation_callback_is_picklable(serializer):
+    results = {}
+    callback = lgb.record_evaluation(eval_result=results)
+    callback_from_disk = pickle_and_unpickle_object(obj=callback, serializer=serializer)
+    assert callback_from_disk.order == 20
+    assert callback_from_disk.before_iteration is False
+    assert callback.eval_result == callback_from_disk.eval_result
+    assert callback.eval_result is results
+
+
+@pytest.mark.parametrize("serializer", SERIALIZERS)
+def test_reset_parameter_callback_is_picklable(serializer):
+    params = {"bagging_fraction": [0.7] * 5 + [0.6] * 5, "feature_fraction": reset_feature_fraction}
+    callback = lgb.reset_parameter(**params)
+    callback_from_disk = pickle_and_unpickle_object(obj=callback, serializer=serializer)
+    assert callback_from_disk.order == 10
+    assert callback_from_disk.before_iteration is True
+    assert callback.kwargs == callback_from_disk.kwargs
+    assert callback.kwargs == params

tests/python_package_test/test_consistency.py

@@ -0,0 +1,143 @@
+# coding: utf-8
+from pathlib import Path
+
+import numpy as np
+from sklearn.datasets import load_svmlight_file
+
+import lightgbm as lgb
+
+EXAMPLES_DIR = Path(__file__).absolute().parents[2] / "examples"
+
+
+class FileLoader:
+    def __init__(self, directory, prefix, config_file="train.conf"):
+        self.directory = directory
+        self.prefix = prefix
+        self.params = {"gpu_use_dp": True}
+        with open(self.directory / config_file, "r") as f:
+            for line in f.readlines():
+                line = line.strip()
+                if line and not line.startswith("#"):
+                    key, value = [token.strip() for token in line.split("=")]
+                    if "early_stopping" not in key:  # disable early_stopping
+                        self.params[key] = value if key not in {"num_trees", "num_threads"} else int(value)
+
+    def load_dataset(self, suffix, is_sparse=False):
+        filename = str(self.path(suffix))
+        if is_sparse:
+            X, Y = load_svmlight_file(filename, dtype=np.float64, zero_based=True)
+            return X, Y, filename
+        else:
+            mat = np.loadtxt(filename, dtype=np.float64)
+            return mat[:, 1:], mat[:, 0], filename
+
+    def load_field(self, suffix):
+        return np.loadtxt(str(self.directory / f"{self.prefix}{suffix}"))
+
+    def load_cpp_result(self, result_file="LightGBM_predict_result.txt"):
+        return np.loadtxt(str(self.directory / result_file))
+
+    def train_predict_check(self, lgb_train, X_test, X_test_fn, sk_pred):
+        params = dict(self.params)
+        params["force_row_wise"] = True
+        gbm = lgb.train(params, lgb_train)
+        y_pred = gbm.predict(X_test)
+        cpp_pred = gbm.predict(X_test_fn)
+        np.testing.assert_allclose(y_pred, cpp_pred)
+        np.testing.assert_allclose(y_pred, sk_pred)
+
+    def file_load_check(self, lgb_train, name):
+        lgb_train_f = lgb.Dataset(self.path(name), params=self.params).construct()
+        for f in ("num_data", "num_feature", "get_label", "get_weight", "get_init_score", "get_group"):
+            a = getattr(lgb_train, f)()
+            b = getattr(lgb_train_f, f)()
+            if a is None and b is None:
+                pass
+            elif a is None:
+                assert np.all(b == 1), f
+            elif isinstance(b, (list, np.ndarray)):
+                np.testing.assert_allclose(a, b)
+            else:
+                assert a == b, f
+
+    def path(self, suffix):
+        return self.directory / f"{self.prefix}{suffix}"
+
+
+def test_binary():
+    fd = FileLoader(EXAMPLES_DIR / "binary_classification", "binary")
+    X_train, y_train, _ = fd.load_dataset(".train")
+    X_test, _, X_test_fn = fd.load_dataset(".test")
+    weight_train = fd.load_field(".train.weight")
+    lgb_train = lgb.Dataset(X_train, y_train, params=fd.params, weight=weight_train)
+    gbm = lgb.LGBMClassifier(**fd.params)
+    gbm.fit(X_train, y_train, sample_weight=weight_train)
+    sk_pred = gbm.predict_proba(X_test)[:, 1]
+    fd.train_predict_check(lgb_train, X_test, X_test_fn, sk_pred)
+    fd.file_load_check(lgb_train, ".train")
+
+
+def test_binary_linear():
+    fd = FileLoader(EXAMPLES_DIR / "binary_classification", "binary", "train_linear.conf")
+    X_train, y_train, _ = fd.load_dataset(".train")
+    X_test, _, X_test_fn = fd.load_dataset(".test")
+    weight_train = fd.load_field(".train.weight")
+    lgb_train = lgb.Dataset(X_train, y_train, params=fd.params, weight=weight_train)
+    gbm = lgb.LGBMClassifier(**fd.params)
+    gbm.fit(X_train, y_train, sample_weight=weight_train)
+    sk_pred = gbm.predict_proba(X_test)[:, 1]
+    fd.train_predict_check(lgb_train, X_test, X_test_fn, sk_pred)
+    fd.file_load_check(lgb_train, ".train")
+
+
+def test_multiclass():
+    fd = FileLoader(EXAMPLES_DIR / "multiclass_classification", "multiclass")
+    X_train, y_train, _ = fd.load_dataset(".train")
+    X_test, _, X_test_fn = fd.load_dataset(".test")
+    lgb_train = lgb.Dataset(X_train, y_train)
+    gbm = lgb.LGBMClassifier(**fd.params)
+    gbm.fit(X_train, y_train)
+    sk_pred = gbm.predict_proba(X_test)
+    fd.train_predict_check(lgb_train, X_test, X_test_fn, sk_pred)
+    fd.file_load_check(lgb_train, ".train")
+
+
+def test_regression():
+    fd = FileLoader(EXAMPLES_DIR / "regression", "regression")
+    X_train, y_train, _ = fd.load_dataset(".train")
+    X_test, _, X_test_fn = fd.load_dataset(".test")
+    init_score_train = fd.load_field(".train.init")
+    lgb_train = lgb.Dataset(X_train, y_train, init_score=init_score_train)
+    gbm = lgb.LGBMRegressor(**fd.params)
+    gbm.fit(X_train, y_train, init_score=init_score_train)
+    sk_pred = gbm.predict(X_test)
+    fd.train_predict_check(lgb_train, X_test, X_test_fn, sk_pred)
+    fd.file_load_check(lgb_train, ".train")
+
+
+def test_lambdarank():
+    fd = FileLoader(EXAMPLES_DIR / "lambdarank", "rank")
+    X_train, y_train, _ = fd.load_dataset(".train", is_sparse=True)
+    X_test, _, X_test_fn = fd.load_dataset(".test", is_sparse=True)
+    group_train = fd.load_field(".train.query")
+    lgb_train = lgb.Dataset(X_train, y_train, group=group_train)
+    params = dict(fd.params)
+    params["force_col_wise"] = True
+    gbm = lgb.LGBMRanker(**params)
+    gbm.fit(X_train, y_train, group=group_train)
+    sk_pred = gbm.predict(X_test)
+    fd.train_predict_check(lgb_train, X_test, X_test_fn, sk_pred)
+    fd.file_load_check(lgb_train, ".train")
+
+
+def test_xendcg():
+    fd = FileLoader(EXAMPLES_DIR / "xendcg", "rank")
+    X_train, y_train, _ = fd.load_dataset(".train", is_sparse=True)
+    X_test, _, X_test_fn = fd.load_dataset(".test", is_sparse=True)
+    group_train = fd.load_field(".train.query")
+    lgb_train = lgb.Dataset(X_train, y_train, group=group_train)
+    gbm = lgb.LGBMRanker(**fd.params)
+    gbm.fit(X_train, y_train, group=group_train)
+    sk_pred = gbm.predict(X_test)
+    fd.train_predict_check(lgb_train, X_test, X_test_fn, sk_pred)
+    fd.file_load_check(lgb_train, ".train")

tests/python_package_test/test_dual.py

@@ -0,0 +1,37 @@
+# coding: utf-8
+"""Tests for dual GPU+CPU support."""
+
+import os
+import platform
+
+import pytest
+from sklearn.metrics import log_loss
+
+import lightgbm as lgb
+
+from .utils import load_breast_cancer
+
+
+@pytest.mark.skipif(
+    os.environ.get("LIGHTGBM_TEST_DUAL_CPU_GPU", "0") != "1",
+    reason="Set LIGHTGBM_TEST_DUAL_CPU_GPU=1 to test using CPU and GPU training from the same package.",
+)
+def test_cpu_and_gpu_work():
+    # If compiled appropriately, the same installation will support both GPU and CPU.
+    X, y = load_breast_cancer(return_X_y=True)
+    data = lgb.Dataset(X, y)
+
+    params_cpu = {"verbosity": -1, "num_leaves": 31, "objective": "binary", "device": "cpu"}
+    cpu_bst = lgb.train(params_cpu, data, num_boost_round=10)
+    cpu_score = log_loss(y, cpu_bst.predict(X))
+
+    params_gpu = params_cpu.copy()
+    params_gpu["device"] = "gpu"
+    # Double-precision floats are only supported on x86_64 with PoCL
+    params_gpu["gpu_use_dp"] = platform.machine() == "x86_64"
+    gpu_bst = lgb.train(params_gpu, data, num_boost_round=10)
+    gpu_score = log_loss(y, gpu_bst.predict(X))
+
+    rel = 1e-6 if params_gpu["gpu_use_dp"] else 1e-4
+    assert cpu_score == pytest.approx(gpu_score, rel=rel)
+    assert gpu_score < 0.242

tests/python_package_test/test_plotting.py

@@ -0,0 +1,631 @@
+# coding: utf-8
+import numpy as np
+import pandas as pd
+import pytest
+from sklearn.model_selection import train_test_split
+
+import lightgbm as lgb
+
+from .utils import load_breast_cancer, make_synthetic_regression
+
+
+@pytest.fixture(scope="function")
+def matplotlib():
+    mpl = pytest.importorskip("matplotlib")
+    # use non-interactive, in-memory renderer
+    mpl.use("Agg")
+    return mpl
+
+
+@pytest.fixture(scope="module")
+def breast_cancer_split():
+    return train_test_split(*load_breast_cancer(return_X_y=True), test_size=0.1, random_state=1)
+
+
+def _categorical_data(category_values_lower_bound, category_values_upper_bound):
+    X, y = load_breast_cancer(return_X_y=True)
+    X_df = pd.DataFrame()
+    rnd = np.random.RandomState(0)
+    n_cat_values = rnd.randint(category_values_lower_bound, category_values_upper_bound, size=X.shape[1])
+    for i in range(X.shape[1]):
+        bins = np.linspace(0, 1, num=n_cat_values[i] + 1)
+        X_df[f"cat_col_{i}"] = pd.qcut(X[:, i], q=bins, labels=range(n_cat_values[i])).as_unordered()
+    return X_df, y
+
+
+@pytest.fixture(scope="module")
+def train_data(breast_cancer_split):
+    X_train, _, y_train, _ = breast_cancer_split
+    return lgb.Dataset(X_train, y_train)
+
+
+@pytest.fixture
+def params():
+    return {"objective": "binary", "verbose": -1, "num_leaves": 3}
+
+
+def test_plot_importance(params, breast_cancer_split, train_data, matplotlib):
+    X_train, _, y_train, _ = breast_cancer_split
+
+    gbm0 = lgb.train(params, train_data, num_boost_round=10)
+    ax0 = lgb.plot_importance(gbm0)
+    assert isinstance(ax0, matplotlib.axes.Axes)
+    assert ax0.get_title() == "Feature importance"
+    assert ax0.get_xlabel() == "Feature importance"
+    assert ax0.get_ylabel() == "Features"
+    assert len(ax0.patches) <= 30
+
+    gbm1 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1)
+    gbm1.fit(X_train, y_train)
+
+    ax1 = lgb.plot_importance(gbm1, color="r", title="t", xlabel="x", ylabel="y")
+    assert isinstance(ax1, matplotlib.axes.Axes)
+    assert ax1.get_title() == "t"
+    assert ax1.get_xlabel() == "x"
+    assert ax1.get_ylabel() == "y"
+    assert len(ax1.patches) <= 30
+    for patch in ax1.patches:
+        assert patch.get_facecolor() == (1.0, 0, 0, 1.0)  # red
+
+    ax2 = lgb.plot_importance(gbm0, color=["r", "y", "g", "b"], title=None, xlabel=None, ylabel=None)
+    assert isinstance(ax2, matplotlib.axes.Axes)
+    assert ax2.get_title() == ""
+    assert ax2.get_xlabel() == ""
+    assert ax2.get_ylabel() == ""
+    assert len(ax2.patches) <= 30
+    assert ax2.patches[0].get_facecolor() == (1.0, 0, 0, 1.0)  # r
+    assert ax2.patches[1].get_facecolor() == (0.75, 0.75, 0, 1.0)  # y
+    assert ax2.patches[2].get_facecolor() == (0, 0.5, 0, 1.0)  # g
+    assert ax2.patches[3].get_facecolor() == (0, 0, 1.0, 1.0)  # b
+
+    ax3 = lgb.plot_importance(
+        gbm0, title="t @importance_type@", xlabel="x @importance_type@", ylabel="y @importance_type@"
+    )
+    assert isinstance(ax3, matplotlib.axes.Axes)
+    assert ax3.get_title() == "t @importance_type@"
+    assert ax3.get_xlabel() == "x split"
+    assert ax3.get_ylabel() == "y @importance_type@"
+    assert len(ax3.patches) <= 30
+
+    ax4 = lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, xlim=(0, 30))
+    assert isinstance(ax4, matplotlib.axes.Axes)
+    assert ax4.get_title() == ""
+    assert ax4.get_xlabel() == ""
+    assert ax4.get_ylabel() == ""
+    assert ax4.get_xlim() == (0, 30)
+    assert len(ax4.patches) <= 30
+
+    with pytest.raises(TypeError, match="xlim must be a tuple of 2 elements."):
+        lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, xlim="not a tuple")
+
+    ax5 = lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, ylim=(0, 30))
+    assert isinstance(ax5, matplotlib.axes.Axes)
+    assert ax5.get_title() == ""
+    assert ax5.get_xlabel() == ""
+    assert ax5.get_ylabel() == ""
+    assert ax5.get_ylim() == (0, 30)
+    assert len(ax5.patches) <= 30
+
+    with pytest.raises(TypeError, match="ylim must be a tuple of 2 elements."):
+        lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, ylim="not a tuple")
+
+    ax6 = lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, figsize=(0, 30))
+    assert isinstance(ax6, matplotlib.axes.Axes)
+    assert ax6.get_title() == ""
+    assert ax6.get_xlabel() == ""
+    assert ax6.get_ylabel() == ""
+    assert list(ax6.get_figure().get_size_inches()) == [0, 30]
+    assert len(ax6.patches) <= 30
+
+    with pytest.raises(TypeError, match="figsize must be a tuple of 2 elements."):
+        lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, figsize="not a tuple")
+
+    # test max_num_features parameter
+    total_features = len(gbm0.feature_importance())
+    assert total_features > 5, "model must have more than 5 features to test max_num_features"
+    ax7 = lgb.plot_importance(gbm0, max_num_features=5)
+    assert isinstance(ax7, matplotlib.axes.Axes)
+    assert len(ax7.patches) == 5
+    # verify the 5 displayed features are the top 5 by importance
+    importance = gbm0.feature_importance()
+    feature_names = gbm0.feature_name()
+    sorted_pairs = sorted(zip(feature_names, importance, strict=True), key=lambda x: x[1])
+    top5_names = [name for name, _ in sorted_pairs[-5:]]
+    displayed_labels = [label.get_text() for label in ax7.get_yticklabels()]
+    assert displayed_labels == top5_names
+
+    gbm2 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1, importance_type="gain")
+    gbm2.fit(X_train, y_train)
+
+    def get_bounds_of_first_patch(axes):
+        return axes.patches[0].get_extents().bounds
+
+    first_bar1 = get_bounds_of_first_patch(lgb.plot_importance(gbm1))
+    first_bar2 = get_bounds_of_first_patch(lgb.plot_importance(gbm1, importance_type="split"))
+    first_bar3 = get_bounds_of_first_patch(lgb.plot_importance(gbm1, importance_type="gain"))
+    first_bar4 = get_bounds_of_first_patch(lgb.plot_importance(gbm2))
+    first_bar5 = get_bounds_of_first_patch(lgb.plot_importance(gbm2, importance_type="split"))
+    first_bar6 = get_bounds_of_first_patch(lgb.plot_importance(gbm2, importance_type="gain"))
+
+    assert first_bar1 == first_bar2
+    assert first_bar1 == first_bar5
+    assert first_bar3 == first_bar4
+    assert first_bar3 == first_bar6
+    assert first_bar1 != first_bar3
+
+
+def test_plot_importance_zero_splits(matplotlib):
+    X, y = load_breast_cancer(return_X_y=True)
+    model = lgb.train(
+        params={
+            "min_data_in_bin": X.shape[0] + 1,
+            "objective": "regression",
+            "verbose": -1,
+        },
+        train_set=lgb.Dataset(X, label=y),
+        num_boost_round=1,
+    )
+    with pytest.raises(ValueError, match="No non-zero feature importances found"):
+        lgb.plot_importance(model)
+    # ignore_zero=False should still produce a valid plot
+    ax = lgb.plot_importance(model, ignore_zero=False)
+    assert isinstance(ax, matplotlib.axes.Axes)
+    assert len(ax.patches) == X.shape[1]
+
+
+def test_plot_split_value_histogram(params, breast_cancer_split, train_data, matplotlib):
+    X_train, _, y_train, _ = breast_cancer_split
+
+    gbm0 = lgb.train(params, train_data, num_boost_round=10)
+    ax0 = lgb.plot_split_value_histogram(gbm0, 27)
+    assert isinstance(ax0, matplotlib.axes.Axes)
+    assert ax0.get_title() == "Split value histogram for feature with index 27"
+    assert ax0.get_xlabel() == "Feature split value"
+    assert ax0.get_ylabel() == "Count"
+    assert len(ax0.patches) <= 2
+
+    gbm1 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1)
+    gbm1.fit(X_train, y_train)
+
+    ax1 = lgb.plot_split_value_histogram(
+        gbm1,
+        gbm1.booster_.feature_name()[27],
+        figsize=(10, 5),
+        title="Histogram for feature @index/name@ @feature@",
+        xlabel="x",
+        ylabel="y",
+        color="r",
+    )
+    assert isinstance(ax1, matplotlib.axes.Axes)
+    title = f"Histogram for feature name {gbm1.booster_.feature_name()[27]}"
+    assert ax1.get_title() == title
+    assert ax1.get_xlabel() == "x"
+    assert ax1.get_ylabel() == "y"
+    assert len(ax1.patches) <= 2
+    for patch in ax1.patches:
+        assert patch.get_facecolor() == (1.0, 0, 0, 1.0)  # red
+
+    ax2 = lgb.plot_split_value_histogram(
+        gbm0, 27, bins=10, color=["r", "y", "g", "b"], title=None, xlabel=None, ylabel=None
+    )
+    assert isinstance(ax2, matplotlib.axes.Axes)
+    assert ax2.get_title() == ""
+    assert ax2.get_xlabel() == ""
+    assert ax2.get_ylabel() == ""
+    assert len(ax2.patches) == 10
+    assert ax2.patches[0].get_facecolor() == (1.0, 0, 0, 1.0)  # r
+    assert ax2.patches[1].get_facecolor() == (0.75, 0.75, 0, 1.0)  # y
+    assert ax2.patches[2].get_facecolor() == (0, 0.5, 0, 1.0)  # g
+    assert ax2.patches[3].get_facecolor() == (0, 0, 1.0, 1.0)  # b
+
+    # test xlim parameter
+    ax3 = lgb.plot_split_value_histogram(gbm0, 27, xlim=(0, 100), title=None, xlabel=None, ylabel=None)
+    assert isinstance(ax3, matplotlib.axes.Axes)
+    assert ax3.get_title() == ""
+    assert ax3.get_xlabel() == ""
+    assert ax3.get_ylabel() == ""
+    assert ax3.get_xlim() == (0, 100)
+
+    with pytest.raises(TypeError, match="xlim must be a tuple of 2 elements."):
+        lgb.plot_split_value_histogram(gbm0, 27, xlim="not a tuple")
+
+    ax4 = lgb.plot_split_value_histogram(gbm0, 27, ylim=(0, 100), title=None, xlabel=None, ylabel=None)
+    assert isinstance(ax4, matplotlib.axes.Axes)
+    assert ax4.get_title() == ""
+    assert ax4.get_xlabel() == ""
+    assert ax4.get_ylabel() == ""
+    assert ax4.get_ylim() == (0, 100)
+
+    with pytest.raises(TypeError, match="ylim must be a tuple of 2 elements."):
+        lgb.plot_split_value_histogram(gbm0, 27, ylim="not a tuple")
+
+    with pytest.raises(
+        ValueError, match="Cannot plot split value histogram, because feature 0 was not used in splitting"
+    ):
+        lgb.plot_split_value_histogram(gbm0, 0)  # was not used in splitting
+
+
+def test_plot_tree(breast_cancer_split, matplotlib):
+    pytest.importorskip("graphviz")
+    X_train, _, y_train, _ = breast_cancer_split
+    gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1)
+    gbm.fit(X_train, y_train)
+
+    with pytest.raises(IndexError, match="tree_index is out of range."):
+        lgb.plot_tree(gbm, tree_index=83)
+
+    ax = lgb.plot_tree(gbm, tree_index=3, figsize=(15, 8), show_info=["split_gain"])
+    assert isinstance(ax, matplotlib.axes.Axes)
+    w, h = ax.axes.get_figure().get_size_inches()
+    assert int(w) == 15
+    assert int(h) == 8
+
+
+def test_create_tree_digraph(tmp_path, breast_cancer_split):
+    graphviz = pytest.importorskip("graphviz")
+    X_train, _, y_train, _ = breast_cancer_split
+
+    constraints = [-1, 1] * int(X_train.shape[1] / 2)
+    gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1, monotone_constraints=constraints)
+    gbm.fit(X_train, y_train)
+
+    with pytest.raises(IndexError, match="tree_index is out of range."):
+        lgb.create_tree_digraph(gbm, tree_index=83)
+
+    graph = lgb.create_tree_digraph(
+        gbm,
+        tree_index=3,
+        show_info=["split_gain", "internal_value", "internal_weight"],
+        name="Tree4",
+        node_attr={"color": "red"},
+        directory=tmp_path,
+    )
+    graph.render(view=False)
+    assert isinstance(graph, graphviz.Digraph)
+    assert graph.name == "Tree4"
+    assert len(graph.node_attr) == 1
+    assert graph.node_attr["color"] == "red"
+    assert len(graph.graph_attr) == 0
+    assert len(graph.edge_attr) == 0
+    graph_body = "".join(graph.body)
+    assert "leaf" in graph_body
+    assert "gain" in graph_body
+    assert "value" in graph_body
+    assert "weight" in graph_body
+    assert "#ffdddd" in graph_body
+    assert "#ddffdd" in graph_body
+    assert "data" not in graph_body
+    assert "count" not in graph_body
+
+
+def test_tree_with_categories_below_max_category_values(tmp_path):
+    graphviz = pytest.importorskip("graphviz")
+    X_train, y_train = _categorical_data(2, 10)
+    params = {
+        "n_estimators": 10,
+        "num_leaves": 3,
+        "min_data_in_bin": 1,
+        "force_col_wise": True,
+        "deterministic": True,
+        "num_threads": 1,
+        "seed": 708,
+        "verbose": -1,
+    }
+    gbm = lgb.LGBMClassifier(**params)
+    gbm.fit(X_train, y_train)
+
+    with pytest.raises(IndexError, match="tree_index is out of range."):
+        lgb.create_tree_digraph(gbm, tree_index=83)
+
+    graph = lgb.create_tree_digraph(
+        gbm,
+        tree_index=3,
+        show_info=["split_gain", "internal_value", "internal_weight"],
+        name="Tree4",
+        node_attr={"color": "red"},
+        max_category_values=10,
+        directory=tmp_path,
+    )
+    graph.render(view=False)
+    assert isinstance(graph, graphviz.Digraph)
+    assert graph.name == "Tree4"
+    assert len(graph.node_attr) == 1
+    assert graph.node_attr["color"] == "red"
+    assert len(graph.graph_attr) == 0
+    assert len(graph.edge_attr) == 0
+    graph_body = "".join(graph.body)
+    assert "leaf" in graph_body
+    assert "gain" in graph_body
+    assert "value" in graph_body
+    assert "weight" in graph_body
+    assert "data" not in graph_body
+    assert "count" not in graph_body
+    assert "||...||" not in graph_body
+
+
+def test_tree_with_categories_above_max_category_values(tmp_path):
+    graphviz = pytest.importorskip("graphviz")
+    X_train, y_train = _categorical_data(20, 30)
+    params = {
+        "n_estimators": 10,
+        "num_leaves": 3,
+        "min_data_in_bin": 1,
+        "force_col_wise": True,
+        "deterministic": True,
+        "num_threads": 1,
+        "seed": 708,
+        "verbose": -1,
+    }
+    gbm = lgb.LGBMClassifier(**params)
+    gbm.fit(X_train, y_train)
+
+    with pytest.raises(IndexError, match="tree_index is out of range."):
+        lgb.create_tree_digraph(gbm, tree_index=83)
+
+    graph = lgb.create_tree_digraph(
+        gbm,
+        tree_index=9,
+        show_info=["split_gain", "internal_value", "internal_weight"],
+        name="Tree4",
+        node_attr={"color": "red"},
+        max_category_values=4,
+        directory=tmp_path,
+    )
+    graph.render(view=False)
+    assert isinstance(graph, graphviz.Digraph)
+    assert graph.name == "Tree4"
+    assert len(graph.node_attr) == 1
+    assert graph.node_attr["color"] == "red"
+    assert len(graph.graph_attr) == 0
+    assert len(graph.edge_attr) == 0
+    graph_body = "".join(graph.body)
+    assert "leaf" in graph_body
+    assert "gain" in graph_body
+    assert "value" in graph_body
+    assert "weight" in graph_body
+    assert "data" not in graph_body
+    assert "count" not in graph_body
+    assert "||...||" in graph_body
+
+
+@pytest.mark.parametrize("use_missing", [True, False])
+@pytest.mark.parametrize("zero_as_missing", [True, False])
+def test_numeric_split_direction(use_missing, zero_as_missing):
+    X, y = make_synthetic_regression()
+    rng = np.random.RandomState(0)
+    zero_mask = rng.rand(X.shape[0]) < 0.05
+    X[zero_mask, :] = 0
+    if use_missing:
+        nan_mask = ~zero_mask & (rng.rand(X.shape[0]) < 0.1)
+        X[nan_mask, :] = np.nan
+    ds = lgb.Dataset(X, y)
+    params = {
+        "num_leaves": 127,
+        "min_child_samples": 1,
+        "use_missing": use_missing,
+        "zero_as_missing": zero_as_missing,
+    }
+    bst = lgb.train(params, ds, num_boost_round=1)
+
+    case_with_zero = X[zero_mask][[0]]
+    expected_leaf_zero = bst.predict(case_with_zero, pred_leaf=True)[0]
+    node = bst.dump_model()["tree_info"][0]["tree_structure"]
+    while "decision_type" in node:
+        direction = lgb.plotting._determine_direction_for_numeric_split(
+            fval=case_with_zero[0][node["split_feature"]],
+            threshold=node["threshold"],
+            missing_type_str=node["missing_type"],
+            default_left=node["default_left"],
+        )
+        node = node["left_child"] if direction == "left" else node["right_child"]
+    assert node["leaf_index"] == expected_leaf_zero
+
+    if use_missing:
+        case_with_nan = X[nan_mask][[0]]
+        expected_leaf_nan = bst.predict(case_with_nan, pred_leaf=True)[0]
+        node = bst.dump_model()["tree_info"][0]["tree_structure"]
+        while "decision_type" in node:
+            direction = lgb.plotting._determine_direction_for_numeric_split(
+                fval=case_with_nan[0][node["split_feature"]],
+                threshold=node["threshold"],
+                missing_type_str=node["missing_type"],
+                default_left=node["default_left"],
+            )
+            node = node["left_child"] if direction == "left" else node["right_child"]
+        assert node["leaf_index"] == expected_leaf_nan
+        if zero_as_missing:
+            # zeros treated as missing -> same leaf as NaN
+            assert expected_leaf_zero == expected_leaf_nan
+        else:
+            # zeros are regular values -> different leaf from NaN
+            assert expected_leaf_zero != expected_leaf_nan
+
+
+def test_example_case_in_tree_digraph():
+    pytest.importorskip("graphviz")
+    rng = np.random.RandomState(0)
+    x1 = rng.rand(100)
+    cat = rng.randint(1, 3, size=x1.size)
+    X = np.vstack([x1, cat]).T
+    y = x1 + 2 * cat
+    feature_name = ["x1", "cat"]
+    ds = lgb.Dataset(X, y, feature_name=feature_name, categorical_feature=["cat"])
+
+    num_round = 3
+    bst = lgb.train({"num_leaves": 7}, ds, num_boost_round=num_round)
+    mod = bst.dump_model()
+    example_case = X[[0]]
+    makes_categorical_splits = False
+    seen_indices = set()
+    for i in range(num_round):
+        graph = lgb.create_tree_digraph(bst, example_case=example_case, tree_index=i)
+        gbody = graph.body
+        node = mod["tree_info"][i]["tree_structure"]
+        while "decision_type" in node:  # iterate through the splits
+            split_index = node["split_index"]
+
+            node_in_graph = [n for n in gbody if f"split{split_index}" in n and "->" not in n]
+            assert len(node_in_graph) == 1
+            seen_indices.add(gbody.index(node_in_graph[0]))
+
+            edge_to_node = [e for e in gbody if f"-> split{split_index}" in e]
+            if node["decision_type"] == "<=":
+                direction = lgb.plotting._determine_direction_for_numeric_split(
+                    fval=example_case[0][node["split_feature"]],
+                    threshold=node["threshold"],
+                    missing_type_str=node["missing_type"],
+                    default_left=node["default_left"],
+                )
+            else:
+                makes_categorical_splits = True
+                direction = lgb.plotting._determine_direction_for_categorical_split(
+                    example_case[0][node["split_feature"]], node["threshold"]
+                )
+            node = node["left_child"] if direction == "left" else node["right_child"]
+            assert "color=blue" in node_in_graph[0]
+            if edge_to_node:
+                assert len(edge_to_node) == 1
+                assert "color=blue" in edge_to_node[0]
+                seen_indices.add(gbody.index(edge_to_node[0]))
+        # we're in a leaf now
+        leaf_index = node["leaf_index"]
+        leaf_in_graph = [n for n in gbody if f"leaf{leaf_index}" in n and "->" not in n]
+        edge_to_leaf = [e for e in gbody if f"-> leaf{leaf_index}" in e]
+        assert len(leaf_in_graph) == 1
+        assert "color=blue" in leaf_in_graph[0]
+        assert len(edge_to_leaf) == 1
+        assert "color=blue" in edge_to_leaf[0]
+        seen_indices.update([gbody.index(leaf_in_graph[0]), gbody.index(edge_to_leaf[0])])
+
+        # check that the rest of the elements have black color
+        remaining_elements = [e for i, e in enumerate(graph.body) if i not in seen_indices and "graph" not in e]
+        assert all("color=black" in e for e in remaining_elements)
+
+        # check that we got to the expected leaf
+        expected_leaf = bst.predict(example_case, start_iteration=i, num_iteration=1, pred_leaf=True)[0]
+        assert leaf_index == expected_leaf
+    assert makes_categorical_splits
+
+
+@pytest.mark.parametrize("input_type", ["array", "dataframe"])
+def test_empty_example_case_on_tree_digraph_raises_error(input_type):
+    pytest.importorskip("graphviz")
+    X, y = make_synthetic_regression()
+    if input_type == "dataframe":
+        pd = pytest.importorskip("pandas")
+        X = pd.DataFrame(X)
+        example_case = pd.DataFrame(X[:0])
+    else:
+        example_case = X[:0]
+    ds = lgb.Dataset(X, y)
+    bst = lgb.train({"num_leaves": 3}, ds, num_boost_round=1)
+    with pytest.raises(ValueError, match="example_case must have a single row."):
+        lgb.create_tree_digraph(bst, tree_index=0, example_case=example_case)
+
+
+def test_plot_metrics(params, breast_cancer_split, train_data, matplotlib):
+    X_train, X_test, y_train, y_test = breast_cancer_split
+    test_data = lgb.Dataset(X_test, y_test, reference=train_data)
+    params.update({"metric": {"binary_logloss", "binary_error"}})
+
+    evals_result0 = {}
+    lgb.train(
+        params,
+        train_data,
+        valid_sets=[train_data, test_data],
+        valid_names=["v1", "v2"],
+        num_boost_round=10,
+        callbacks=[lgb.record_evaluation(evals_result0)],
+    )
+    with pytest.warns(UserWarning, match="More than one metric available, picking one to plot."):
+        ax0 = lgb.plot_metric(evals_result0)
+    assert isinstance(ax0, matplotlib.axes.Axes)
+    assert ax0.get_title() == "Metric during training"
+    assert ax0.get_xlabel() == "Iterations"
+    assert ax0.get_ylabel() in {"binary_logloss", "binary_error"}
+    legend_items = ax0.get_legend().get_texts()
+    assert len(legend_items) == 2
+    assert legend_items[0].get_text() == "v1"
+    assert legend_items[1].get_text() == "v2"
+
+    ax1 = lgb.plot_metric(evals_result0, metric="binary_error")
+    assert isinstance(ax1, matplotlib.axes.Axes)
+    assert ax1.get_title() == "Metric during training"
+    assert ax1.get_xlabel() == "Iterations"
+    assert ax1.get_ylabel() == "binary_error"
+    legend_items = ax1.get_legend().get_texts()
+    assert len(legend_items) == 2
+    assert legend_items[0].get_text() == "v1"
+    assert legend_items[1].get_text() == "v2"
+
+    ax2 = lgb.plot_metric(evals_result0, metric="binary_logloss", dataset_names=["v2"])
+    assert isinstance(ax2, matplotlib.axes.Axes)
+    assert ax2.get_title() == "Metric during training"
+    assert ax2.get_xlabel() == "Iterations"
+    assert ax2.get_ylabel() == "binary_logloss"
+    legend_items = ax2.get_legend().get_texts()
+    assert len(legend_items) == 1
+    assert legend_items[0].get_text() == "v2"
+
+    ax3 = lgb.plot_metric(
+        evals_result0,
+        metric="binary_logloss",
+        dataset_names=["v1"],
+        title="Metric @metric@",
+        xlabel="Iterations @metric@",
+        ylabel='Value of "@metric@"',
+        figsize=(5, 5),
+        dpi=600,
+        grid=False,
+    )
+    assert isinstance(ax3, matplotlib.axes.Axes)
+    assert ax3.get_title() == "Metric @metric@"
+    assert ax3.get_xlabel() == "Iterations @metric@"
+    assert ax3.get_ylabel() == 'Value of "binary_logloss"'
+    legend_items = ax3.get_legend().get_texts()
+    assert len(legend_items) == 1
+    assert legend_items[0].get_text() == "v1"
+    assert ax3.get_figure().get_figheight() == 5
+    assert ax3.get_figure().get_figwidth() == 5
+    assert ax3.get_figure().get_dpi() == 600
+    for grid_line in ax3.get_xgridlines():
+        assert not grid_line.get_visible()
+    for grid_line in ax3.get_ygridlines():
+        assert not grid_line.get_visible()
+
+    evals_result1 = {}
+    lgb.train(params, train_data, num_boost_round=10, callbacks=[lgb.record_evaluation(evals_result1)])
+    with pytest.raises(ValueError, match="eval results cannot be empty."):
+        lgb.plot_metric(evals_result1)
+
+    gbm2 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1)
+    gbm2.fit(X_train, y_train, eval_set=[(X_test, y_test)])
+    ax4 = lgb.plot_metric(gbm2, title=None, xlabel=None, ylabel=None)
+    assert isinstance(ax4, matplotlib.axes.Axes)
+    assert ax4.get_title() == ""
+    assert ax4.get_xlabel() == ""
+    assert ax4.get_ylabel() == ""
+    legend_items = ax4.get_legend().get_texts()
+    assert len(legend_items) == 1
+    assert legend_items[0].get_text() == "valid_0"
+
+    # test xlim parameter
+    ax5 = lgb.plot_metric(evals_result0, metric="binary_logloss", xlim=(0, 15), title=None, xlabel=None, ylabel=None)
+    assert isinstance(ax5, matplotlib.axes.Axes)
+    assert ax5.get_title() == ""
+    assert ax5.get_xlabel() == ""
+    assert ax5.get_ylabel() == ""
+    assert ax5.get_xlim() == (0, 15)
+
+    with pytest.raises(TypeError, match="xlim must be a tuple of 2 elements."):
+        lgb.plot_metric(evals_result0, metric="binary_logloss", xlim="not a tuple")
+
+    ax6 = lgb.plot_metric(evals_result0, metric="binary_logloss", ylim=(0, 15), title=None, xlabel=None, ylabel=None)
+    assert isinstance(ax6, matplotlib.axes.Axes)
+    assert ax6.get_title() == ""
+    assert ax6.get_xlabel() == ""
+    assert ax6.get_ylabel() == ""
+    assert ax6.get_ylim() == (0, 15)
+
+    with pytest.raises(TypeError, match="ylim must be a tuple of 2 elements."):
+        lgb.plot_metric(evals_result0, metric="binary_logloss", ylim="not a tuple")

tests/python_package_test/test_polars.py

@@ -0,0 +1,413 @@
+# coding: utf-8
+import filecmp
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import numpy as np
+import pytest
+
+import lightgbm as lgb
+
+from .utils import np_assert_array_equal
+
+pl = pytest.importorskip("polars")
+
+
+# ----------------------------------------------------------------------------------------------- #
+#                                            UTILITIES                                            #
+# ----------------------------------------------------------------------------------------------- #
+
+_INTEGER_TYPES = [pl.Int8, pl.Int16, pl.Int32, pl.Int64, pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64]
+_FLOAT_TYPES = [pl.Float32, pl.Float64]
+
+
+def generate_simple_polars_frame() -> pl.DataFrame:
+    values = [1, 2, 3, 4, 5]
+    bool_values = [True, True, False, False, True]
+    columns = {f"col_{i}": pl.Series(values, dtype=dtype) for i, dtype in enumerate(_INTEGER_TYPES + _FLOAT_TYPES)}
+    columns[f"col_{len(columns)}"] = pl.Series(bool_values, dtype=pl.Boolean)
+    return pl.DataFrame(columns)
+
+
+def generate_nullable_polars_frame(dtype: Any) -> pl.DataFrame:
+    return pl.DataFrame(
+        {
+            "col_0": pl.Series([1, None, 3, 4, 5], dtype=dtype),
+            "col_1": pl.Series([None, 2, 3, 4, 5], dtype=dtype),
+            "col_2": pl.Series([1, 2, 3, 4, None], dtype=dtype),
+            "col_3": pl.Series([None, None, None, None, None], dtype=dtype),
+        }
+    )
+
+
+def generate_dummy_polars_frame() -> pl.DataFrame:
+    return pl.DataFrame(
+        {
+            "a": pl.Series([1, 2, 3, 4, 5], dtype=pl.UInt8),
+            "b": pl.Series([0.5, 0.6, 0.1, 0.8, 1.5], dtype=pl.Float32),
+        }
+    )
+
+
+def generate_random_polars_frame(
+    num_columns: int,
+    num_datapoints: int,
+    seed: int,
+    generate_nulls: bool = True,
+    values: Optional[np.ndarray] = None,
+) -> pl.DataFrame:
+    return pl.DataFrame(
+        {
+            f"col_{i}": generate_random_polars_series(
+                num_datapoints, seed + i, generate_nulls=generate_nulls, values=values
+            )
+            for i in range(num_columns)
+        }
+    )
+
+
+def generate_random_polars_series(
+    num_datapoints: int,
+    seed: int,
+    generate_nulls: bool = True,
+    values: Optional[np.ndarray] = None,
+) -> pl.Series:
+    generator = np.random.default_rng(seed)
+    data = (
+        generator.standard_normal(num_datapoints).astype(np.float32)
+        if values is None
+        else generator.choice(values, size=num_datapoints, replace=True)
+    )
+    series = pl.Series("col", data, dtype=pl.Float32)
+    if generate_nulls:
+        indices = generator.choice(len(data), size=num_datapoints // 10)
+        series = series.scatter(indices, None)
+    return series
+
+
+def dummy_dataset_params() -> Dict[str, Any]:
+    return {
+        "min_data_in_bin": 1,
+        "min_data_in_leaf": 1,
+    }
+
+
+# ----------------------------------------------------------------------------------------------- #
+#                                            UNIT TESTS                                           #
+# ----------------------------------------------------------------------------------------------- #
+
+# ------------------------------------------- DATASET ------------------------------------------- #
+
+
+def assert_datasets_equal(tmp_path: Path, lhs: lgb.Dataset, rhs: lgb.Dataset):
+    lhs._dump_text(tmp_path / "polars.txt")
+    rhs._dump_text(tmp_path / "pandas.txt")
+    assert filecmp.cmp(tmp_path / "polars.txt", tmp_path / "pandas.txt")
+
+
+@pytest.mark.parametrize(
+    ("polars_frame_fn", "dataset_params"),
+    [  # Use lambda functions here to minimize memory consumption
+        (generate_simple_polars_frame, dummy_dataset_params()),
+        (generate_dummy_polars_frame, dummy_dataset_params()),
+        (lambda: generate_nullable_polars_frame(pl.Float32), dummy_dataset_params()),
+        (lambda: generate_nullable_polars_frame(pl.Int32), dummy_dataset_params()),
+        (lambda: generate_random_polars_frame(3, 1000, 42), {}),
+        (lambda: generate_random_polars_frame(100, 10000, 43), {}),
+    ],
+)
+def test_dataset_construct_fuzzy(tmp_path, polars_frame_fn, dataset_params):
+    polars_frame = polars_frame_fn()
+
+    polars_dataset = lgb.Dataset(polars_frame, params=dataset_params)
+    polars_dataset.construct()
+
+    pandas_dataset = lgb.Dataset(polars_frame.to_pandas(), params=dataset_params)
+    pandas_dataset.construct()
+
+    assert_datasets_equal(tmp_path, polars_dataset, pandas_dataset)
+
+
+def test_dataset_construct_fuzzy_boolean(tmp_path):
+    boolean_data = generate_random_polars_frame(10, 10000, 42, generate_nulls=False, values=np.array([True, False]))
+    float_data = boolean_data.cast(pl.Float32)
+
+    polars_dataset = lgb.Dataset(boolean_data)
+    polars_dataset.construct()
+
+    pandas_dataset = lgb.Dataset(float_data.to_pandas())
+    pandas_dataset.construct()
+
+    assert_datasets_equal(tmp_path, polars_dataset, pandas_dataset)
+
+
+# -------------------------------------------- FIELDS ------------------------------------------- #
+
+
+def test_dataset_construct_fields_fuzzy():
+    polars_frame = generate_random_polars_frame(3, 1000, 42)
+    polars_labels = generate_random_polars_series(1000, 42, generate_nulls=False)
+    polars_weights = generate_random_polars_series(1000, 42, generate_nulls=False)
+    polars_groups = pl.Series("group", [300, 400, 50, 250], dtype=pl.Int32)
+
+    polars_dataset = lgb.Dataset(polars_frame, label=polars_labels, weight=polars_weights, group=polars_groups)
+    polars_dataset.construct()
+
+    pandas_dataset = lgb.Dataset(
+        polars_frame.to_pandas(),
+        label=polars_labels.to_numpy(),
+        weight=polars_weights.to_numpy(),
+        group=polars_groups.to_numpy(),
+    )
+    pandas_dataset.construct()
+
+    # Check for equality
+    for field in ("label", "weight", "group"):
+        np_assert_array_equal(polars_dataset.get_field(field), pandas_dataset.get_field(field), strict=True)
+    np_assert_array_equal(polars_dataset.get_label(), pandas_dataset.get_label(), strict=True)
+    np_assert_array_equal(polars_dataset.get_weight(), pandas_dataset.get_weight(), strict=True)
+
+
+# -------------------------------------------- LABELS ------------------------------------------- #
+
+
+@pytest.mark.parametrize("polars_type", _INTEGER_TYPES + _FLOAT_TYPES)
+def test_dataset_construct_labels(polars_type):
+    data = generate_dummy_polars_frame()
+    labels = pl.Series("label", [0, 1, 0, 0, 1], dtype=polars_type)
+    dataset = lgb.Dataset(data, label=labels, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 0, 0, 1], dtype=np.float32)
+    np_assert_array_equal(expected, dataset.get_label(), strict=True)
+
+
+def test_dataset_construct_labels_boolean():
+    data = generate_dummy_polars_frame()
+    labels = pl.Series("label", [False, True, False, False, True], dtype=pl.Boolean)
+    dataset = lgb.Dataset(data, label=labels, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 0, 0, 1], dtype=np.float32)
+    np_assert_array_equal(expected, dataset.get_label(), strict=True)
+
+
+# ------------------------------------------- WEIGHTS ------------------------------------------- #
+
+
+def test_dataset_construct_weights_none():
+    data = generate_dummy_polars_frame()
+    weight = pl.Series("weight", [1, 1, 1, 1, 1], dtype=pl.Float32)
+    dataset = lgb.Dataset(data, weight=weight, params=dummy_dataset_params())
+    dataset.construct()
+    assert dataset.get_weight() is None
+    assert dataset.get_field("weight") is None
+
+
+@pytest.mark.parametrize("polars_type", _FLOAT_TYPES)
+def test_dataset_construct_weights(polars_type):
+    data = generate_dummy_polars_frame()
+    weights = pl.Series("weight", [3, 0.7, 1.5, 0.5, 0.1], dtype=polars_type)
+    dataset = lgb.Dataset(data, weight=weights, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([3, 0.7, 1.5, 0.5, 0.1], dtype=np.float32)
+    np_assert_array_equal(expected, dataset.get_weight(), strict=True)
+
+
+# -------------------------------------------- GROUPS ------------------------------------------- #
+
+
+@pytest.mark.parametrize("polars_type", _INTEGER_TYPES)
+def test_dataset_construct_groups(polars_type):
+    data = generate_dummy_polars_frame()
+    groups = pl.Series("group", [2, 3], dtype=polars_type)
+    dataset = lgb.Dataset(data, group=groups, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 2, 5], dtype=np.int32)
+    np_assert_array_equal(expected, dataset.get_field("group"), strict=True)
+
+
+# ----------------------------------------- INIT SCORES ----------------------------------------- #
+
+
+@pytest.mark.parametrize("polars_type", _INTEGER_TYPES + _FLOAT_TYPES)
+def test_dataset_construct_init_scores_array(polars_type):
+    data = generate_dummy_polars_frame()
+    init_scores = pl.Series("init_score", [0, 1, 2, 3, 3], dtype=polars_type)
+    dataset = lgb.Dataset(data, init_score=init_scores, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 2, 3, 3], dtype=np.float64)
+    np_assert_array_equal(expected, dataset.get_init_score(), strict=True)
+
+
+def test_dataset_construct_init_scores_table():
+    data = generate_dummy_polars_frame()
+    init_scores = pl.DataFrame(
+        {
+            "a": generate_random_polars_series(5, seed=1, generate_nulls=False),
+            "b": generate_random_polars_series(5, seed=2, generate_nulls=False),
+            "c": generate_random_polars_series(5, seed=3, generate_nulls=False),
+        }
+    )
+    dataset = lgb.Dataset(data, init_score=init_scores, params=dummy_dataset_params())
+    dataset.construct()
+
+    actual = dataset.get_init_score()
+    expected = init_scores.to_numpy().astype(np.float64)
+    np_assert_array_equal(expected, actual, strict=True)
+
+
+# ------------------------------------------ PREDICTION ----------------------------------------- #
+
+
+def assert_equal_predict_polars_pandas(booster: lgb.Booster, data: pl.DataFrame):
+    pandas_data = data.to_pandas()
+
+    p_polars = booster.predict(data)
+    p_pandas = booster.predict(pandas_data)
+    np_assert_array_equal(p_polars, p_pandas, strict=True)
+
+    p_raw_polars = booster.predict(data, raw_score=True)
+    p_raw_pandas = booster.predict(pandas_data, raw_score=True)
+    np_assert_array_equal(p_raw_polars, p_raw_pandas, strict=True)
+
+    p_leaf_polars = booster.predict(data, pred_leaf=True)
+    p_leaf_pandas = booster.predict(pandas_data, pred_leaf=True)
+    np_assert_array_equal(p_leaf_polars, p_leaf_pandas, strict=True)
+
+    p_pred_contrib_polars = booster.predict(data, pred_contrib=True)
+    p_pred_contrib_pandas = booster.predict(pandas_data, pred_contrib=True)
+    np_assert_array_equal(p_pred_contrib_polars, p_pred_contrib_pandas, strict=True)
+
+    p_first_iter_polars = booster.predict(data, start_iteration=0, num_iteration=1, raw_score=True)
+    p_first_iter_pandas = booster.predict(pandas_data, start_iteration=0, num_iteration=1, raw_score=True)
+    np_assert_array_equal(p_first_iter_polars, p_first_iter_pandas, strict=True)
+
+
+def test_predict_regression():
+    data_float = generate_random_polars_frame(10, 10000, 42)
+    data_bool = generate_random_polars_frame(1, 10000, 42, generate_nulls=False, values=np.array([True, False]))
+    data = data_float.with_columns(data_bool["col_0"].alias("col_bool"))
+
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_polars_series(10000, 43, generate_nulls=False),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "regression", "num_leaves": 7},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_polars_pandas(booster, data)
+
+
+def test_predict_binary_classification():
+    data = generate_random_polars_frame(10, 10000, 42)
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_polars_series(10000, 43, generate_nulls=False, values=np.arange(2)),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "binary", "num_leaves": 7},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_polars_pandas(booster, data)
+
+
+def test_predict_multiclass_classification():
+    data = generate_random_polars_frame(10, 10000, 42)
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_polars_series(10000, 43, generate_nulls=False, values=np.arange(5)),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "multiclass", "num_leaves": 7, "num_class": 5},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_polars_pandas(booster, data)
+
+
+def test_predict_ranking():
+    data = generate_random_polars_frame(10, 10000, 42)
+    dataset = lgb.Dataset(
+        data,
+        label=generate_random_polars_series(10000, 43, generate_nulls=False, values=np.arange(4)),
+        group=np.array([1000, 2000, 3000, 4000]),
+        params=dummy_dataset_params(),
+    )
+    booster = lgb.train(
+        {"objective": "lambdarank", "num_leaves": 7},
+        dataset,
+        num_boost_round=5,
+    )
+    assert_equal_predict_polars_pandas(booster, data)
+
+
+def test_polars_feature_name_auto():
+    data = generate_dummy_polars_frame()
+    dataset = lgb.Dataset(
+        data,
+        label=pl.Series("label", [0, 1, 0, 0, 1]),
+        params=dummy_dataset_params(),
+        categorical_feature=["a"],
+    )
+    booster = lgb.train({"num_leaves": 7}, dataset, num_boost_round=5)
+    assert booster.feature_name() == ["a", "b"]
+
+
+def test_polars_feature_name_manual():
+    data = generate_dummy_polars_frame()
+    dataset = lgb.Dataset(
+        data,
+        label=pl.Series("label", [0, 1, 0, 0, 1]),
+        params=dummy_dataset_params(),
+        feature_name=["c", "d"],
+        categorical_feature=["c"],
+    )
+    booster = lgb.train({"num_leaves": 7}, dataset, num_boost_round=5)
+    assert booster.feature_name() == ["c", "d"]
+
+
+def test_get_data_polars_frame():
+    from polars.testing import assert_frame_equal  # noqa: PLC0415
+
+    original_frame = generate_simple_polars_frame()
+    dataset = lgb.Dataset(original_frame, free_raw_data=False)
+    dataset.construct()
+
+    returned_data = dataset.get_data()
+    assert isinstance(returned_data, pl.DataFrame)
+    assert returned_data.schema == original_frame.schema
+    assert returned_data.shape == original_frame.shape
+    assert_frame_equal(returned_data, original_frame)
+
+
+def test_get_data_polars_frame_subset(rng):
+    from polars.testing import assert_frame_equal  # noqa: PLC0415
+
+    original_frame = generate_random_polars_frame(num_columns=3, num_datapoints=1000, seed=42)
+    dataset = lgb.Dataset(original_frame, free_raw_data=False)
+    dataset.construct()
+
+    subset_size = 100
+    used_indices = rng.choice(a=original_frame.shape[0], size=subset_size, replace=False)
+    used_indices = sorted(used_indices)
+
+    subset_dataset = dataset.subset(used_indices).construct()
+    expected_subset = original_frame[used_indices]
+    subset_data = subset_dataset.get_data()
+
+    assert isinstance(subset_data, pl.DataFrame)
+    assert subset_data.schema == expected_subset.schema
+    assert subset_data.shape == expected_subset.shape
+    assert len(subset_data) == len(used_indices)
+    assert subset_data.shape == (subset_size, 3)
+    assert_frame_equal(subset_data, expected_subset)

tests/python_package_test/test_utilities.py

@@ -0,0 +1,160 @@
+# coding: utf-8
+import logging
+
+import numpy as np
+import pytest
+
+import lightgbm as lgb
+
+
+def test_register_logger(tmp_path):
+    logger = logging.getLogger("LightGBM")
+    logger.setLevel(logging.DEBUG)
+    formatter = logging.Formatter("%(levelname)s | %(message)s")
+    log_filename = tmp_path / "LightGBM_test_logger.log"
+    file_handler = logging.FileHandler(log_filename, mode="w", encoding="utf-8")
+    file_handler.setLevel(logging.DEBUG)
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    def dummy_metric(_, __):
+        logger.debug("In dummy_metric")
+        return "dummy_metric", 1, True
+
+    lgb.register_logger(logger)
+
+    X = np.array([[1, 2, 3], [1, 2, 4], [1, 2, 4], [1, 2, 3]], dtype=np.float32)
+    y = np.array([0, 1, 1, 0])
+    lgb_train = lgb.Dataset(X, y, categorical_feature=[1])
+    lgb_valid = lgb.Dataset(X, y, categorical_feature=[1])  # different object for early-stopping
+
+    eval_records = {}
+    callbacks = [lgb.record_evaluation(eval_records), lgb.log_evaluation(2), lgb.early_stopping(10)]
+    lgb.train(
+        {"objective": "binary", "metric": ["auc", "binary_error"], "verbose": 1},
+        lgb_train,
+        num_boost_round=10,
+        feval=dummy_metric,
+        valid_sets=[lgb_valid],
+        callbacks=callbacks,
+    )
+
+    lgb.plot_metric(eval_records)
+
+    expected_log = r"""
+INFO | [LightGBM] [Warning] There are no meaningful features which satisfy the provided configuration. Decreasing Dataset parameters min_data_in_bin or min_data_in_leaf and re-constructing Dataset might resolve this warning.
+INFO | [LightGBM] [Info] Number of positive: 2, number of negative: 2
+INFO | [LightGBM] [Info] Total Bins 0
+INFO | [LightGBM] [Info] Number of data points in the train set: 4, number of used features: 0
+INFO | [LightGBM] [Info] [binary:BoostFromScore]: pavg=0.500000 -> initscore=0.000000
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | Training until validation scores don't improve for 10 rounds
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [2]	valid_0's auc: 0.5	valid_0's binary_error: 0.5	valid_0's dummy_metric: 1
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [4]	valid_0's auc: 0.5	valid_0's binary_error: 0.5	valid_0's dummy_metric: 1
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [6]	valid_0's auc: 0.5	valid_0's binary_error: 0.5	valid_0's dummy_metric: 1
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [8]	valid_0's auc: 0.5	valid_0's binary_error: 0.5	valid_0's dummy_metric: 1
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [LightGBM] [Warning] Stopped training because there are no more leaves that meet the split requirements
+DEBUG | In dummy_metric
+INFO | [10]	valid_0's auc: 0.5	valid_0's binary_error: 0.5	valid_0's dummy_metric: 1
+INFO | Did not meet early stopping. Best iteration is:
+[1]	valid_0's auc: 0.5	valid_0's binary_error: 0.5	valid_0's dummy_metric: 1
+WARNING | More than one metric available, picking one to plot.
+""".strip()
+
+    gpu_lines = [
+        "INFO | [LightGBM] [Info] This is the GPU trainer",
+        "INFO | [LightGBM] [Info] Using GPU Device:",
+        "INFO | [LightGBM] [Info] Compiling OpenCL Kernel with 16 bins...",
+        "INFO | [LightGBM] [Info] GPU programs have been built",
+        "INFO | [LightGBM] [Warning] GPU acceleration is disabled because no non-trivial dense features can be found",
+        "INFO | [LightGBM] [Warning] Using sparse features with CUDA is currently not supported.",
+        "INFO | [LightGBM] [Warning] CUDA currently requires double precision calculations.",
+        "INFO | [LightGBM] [Info] LightGBM using CUDA trainer with DP float!!",
+    ]
+    cuda_lines = [
+        "INFO | [LightGBM] [Warning] Metric auc is not implemented in cuda version. Fall back to evaluation on CPU.",
+        "INFO | [LightGBM] [Warning] Metric binary_error is not implemented in cuda version. Fall back to evaluation on CPU.",
+    ]
+    with open(log_filename, "rt", encoding="utf-8") as f:
+        actual_log = f.read().strip()
+        actual_log_wo_gpu_stuff = []
+        for line in actual_log.split("\n"):
+            if not any(line.startswith(gpu_or_cuda_line) for gpu_or_cuda_line in gpu_lines + cuda_lines):
+                actual_log_wo_gpu_stuff.append(line)
+
+    assert "\n".join(actual_log_wo_gpu_stuff) == expected_log
+
+
+def test_register_invalid_logger():
+    class LoggerWithoutInfoMethod:
+        def warning(self, msg: str) -> None:
+            print(msg)
+
+    class LoggerWithoutWarningMethod:
+        def info(self, msg: str) -> None:
+            print(msg)
+
+    class LoggerWithAttributeNotCallable:
+        def __init__(self):
+            self.info = 1
+            self.warning = 2
+
+    expected_error_message = "Logger must provide 'info' and 'warning' method"
+
+    with pytest.raises(TypeError, match=expected_error_message):
+        lgb.register_logger(LoggerWithoutInfoMethod())
+
+    with pytest.raises(TypeError, match=expected_error_message):
+        lgb.register_logger(LoggerWithoutWarningMethod())
+
+    with pytest.raises(TypeError, match=expected_error_message):
+        lgb.register_logger(LoggerWithAttributeNotCallable())
+
+
+def test_register_custom_logger():
+    logged_messages = []
+
+    class CustomLogger:
+        def custom_info(self, msg: str) -> None:
+            logged_messages.append(msg)
+
+        def custom_warning(self, msg: str) -> None:
+            logged_messages.append(msg)
+
+    custom_logger = CustomLogger()
+    lgb.register_logger(custom_logger, info_method_name="custom_info", warning_method_name="custom_warning")
+
+    lgb.basic._log_info("info message")
+    lgb.basic._log_warning("warning message")
+
+    expected_log = ["info message", "warning message"]
+    assert logged_messages == expected_log
+
+    logged_messages = []
+    X = np.array([[1, 2, 3], [1, 2, 4], [1, 2, 4], [1, 2, 3]], dtype=np.float32)
+    y = np.array([0, 1, 1, 0])
+    lgb_data = lgb.Dataset(X, y, categorical_feature=[1])
+    lgb.train(
+        {"objective": "binary", "metric": "auc"},
+        lgb_data,
+        num_boost_round=10,
+        valid_sets=[lgb_data],
+    )
+    assert logged_messages, "custom logger was not called"

tests/python_package_test/utils.py

@@ -0,0 +1,279 @@
+# coding: utf-8
+import os
+import pickle
+from functools import lru_cache
+from inspect import getfullargspec
+
+import cloudpickle
+import joblib
+import numpy as np
+import sklearn.datasets
+from sklearn.utils import check_random_state
+
+import lightgbm as lgb
+
+SERIALIZERS = ["pickle", "joblib", "cloudpickle"]
+
+
+@lru_cache(maxsize=None)
+def load_breast_cancer(**kwargs):
+    return sklearn.datasets.load_breast_cancer(**kwargs)
+
+
+@lru_cache(maxsize=None)
+def load_digits(**kwargs):
+    return sklearn.datasets.load_digits(**kwargs)
+
+
+@lru_cache(maxsize=None)
+def load_iris(**kwargs):
+    return sklearn.datasets.load_iris(**kwargs)
+
+
+@lru_cache(maxsize=None)
+def load_linnerud(**kwargs):
+    return sklearn.datasets.load_linnerud(**kwargs)
+
+
+def make_ranking(
+    *, n_samples=100, n_features=20, n_informative=5, gmax=2, group=None, random_gs=False, avg_gs=10, random_state=0
+):
+    """Generate a learning-to-rank dataset - feature vectors grouped together with
+    integer-valued graded relevance scores. Replace this with a sklearn.datasets function
+    if ranking objective becomes supported in sklearn.datasets module.
+
+    Parameters
+    ----------
+    n_samples : int, optional (default=100)
+        Total number of documents (records) in the dataset.
+    n_features : int, optional (default=20)
+        Total number of features in the dataset.
+    n_informative : int, optional (default=5)
+        Number of features that are "informative" for ranking, as they are bias + beta * y
+        where bias and beta are standard normal variates. If this is greater than n_features, the dataset will have
+        n_features features, all will be informative.
+    gmax : int, optional (default=2)
+        Maximum graded relevance value for creating relevance/target vector. If you set this to 2, for example, all
+        documents in a group will have relevance scores of either 0, 1, or 2.
+    group : array-like, optional (default=None)
+        1-d array or list of group sizes. When `group` is specified, this overrides n_samples, random_gs, and
+        avg_gs by simply creating groups with sizes group[0], ..., group[-1].
+    random_gs : bool, optional (default=False)
+        True will make group sizes ~ Poisson(avg_gs), False will make group sizes == avg_gs.
+    avg_gs : int, optional (default=10)
+        Average number of documents (records) in each group.
+    random_state : int, optional (default=0)
+        Random seed.
+
+    Returns
+    -------
+    X : 2-d np.ndarray of shape = [n_samples (or np.sum(group)), n_features]
+        Input feature matrix for ranking objective.
+    y : 1-d np.array of shape = [n_samples (or np.sum(group))]
+        Integer-graded relevance scores.
+    group_ids : 1-d np.array of shape = [n_samples (or np.sum(group))]
+        Array of group ids, each value indicates to which group each record belongs.
+    """
+    rnd_generator = check_random_state(random_state)
+
+    y_vec, group_id_vec = np.empty((0,), dtype=int), np.empty((0,), dtype=int)
+    gid = 0
+
+    # build target, group ID vectors.
+    relvalues = range(gmax + 1)
+
+    # build y/target and group-id vectors with user-specified group sizes.
+    if group is not None and hasattr(group, "__len__"):
+        n_samples = np.sum(group)
+
+        for i, gsize in enumerate(group):
+            y_vec = np.concatenate((y_vec, rnd_generator.choice(relvalues, size=gsize, replace=True)))
+            group_id_vec = np.concatenate((group_id_vec, [i] * gsize))
+
+    # build y/target and group-id vectors according to n_samples, avg_gs, and random_gs.
+    else:
+        while len(y_vec) < n_samples:
+            gsize = avg_gs if not random_gs else rnd_generator.poisson(avg_gs)
+
+            # groups should contain > 1 element for pairwise learning objective.
+            if gsize < 1:
+                continue
+
+            y_vec = np.append(y_vec, rnd_generator.choice(relvalues, size=gsize, replace=True))
+            group_id_vec = np.append(group_id_vec, [gid] * gsize)
+            gid += 1
+
+        y_vec, group_id_vec = y_vec[:n_samples], group_id_vec[:n_samples]
+
+    # build feature data, X. Transform first few into informative features.
+    n_informative = max(min(n_features, n_informative), 0)
+    X = rnd_generator.uniform(size=(n_samples, n_features))
+
+    for j in range(n_informative):
+        bias, coef = rnd_generator.normal(size=2)
+        X[:, j] = bias + coef * y_vec
+
+    return X, y_vec, group_id_vec
+
+
+@lru_cache(maxsize=None)
+def make_synthetic_regression(*, n_samples=100, n_features=4, n_informative=2, random_state=42):
+    return sklearn.datasets.make_regression(
+        n_samples=n_samples, n_features=n_features, n_informative=n_informative, random_state=random_state
+    )
+
+
+def dummy_obj(preds, train_data):
+    return np.ones(preds.shape), np.ones(preds.shape)
+
+
+def mse_obj(y_pred, dtrain):
+    y_true = dtrain.get_label()
+    grad = y_pred - y_true
+    hess = np.ones(len(grad))
+    return grad, hess
+
+
+def softmax(x):
+    row_wise_max = np.max(x, axis=1).reshape(-1, 1)
+    exp_x = np.exp(x - row_wise_max)
+    return exp_x / np.sum(exp_x, axis=1).reshape(-1, 1)
+
+
+def logistic_sigmoid(x):
+    return 1.0 / (1.0 + np.exp(-x))
+
+
+def sklearn_multiclass_custom_objective(y_true, y_pred, weight=None):
+    num_rows, num_class = y_pred.shape
+    prob = softmax(y_pred)
+    grad_update = np.zeros_like(prob)
+    grad_update[np.arange(num_rows), y_true.astype(np.int32)] = -1.0
+    grad = prob + grad_update
+    factor = num_class / (num_class - 1)
+    hess = factor * prob * (1 - prob)
+    if weight is not None:
+        weight2d = weight.reshape(-1, 1)
+        grad *= weight2d
+        hess *= weight2d
+    return grad, hess
+
+
+def pickle_obj(obj, filepath, serializer):
+    if serializer == "pickle":
+        with open(filepath, "wb") as f:
+            pickle.dump(obj, f)
+    elif serializer == "joblib":
+        joblib.dump(obj, filepath)
+    elif serializer == "cloudpickle":
+        with open(filepath, "wb") as f:
+            cloudpickle.dump(obj, f)
+    else:
+        raise ValueError(f"Unrecognized serializer type: {serializer}")
+
+
+def unpickle_obj(filepath, serializer):
+    if serializer == "pickle":
+        with open(filepath, "rb") as f:
+            return pickle.load(f)
+    elif serializer == "joblib":
+        return joblib.load(filepath)
+    elif serializer == "cloudpickle":
+        with open(filepath, "rb") as f:
+            return cloudpickle.load(f)
+    else:
+        raise ValueError(f"Unrecognized serializer type: {serializer}")
+
+
+def pickle_and_unpickle_object(obj, serializer):
+    with lgb.basic._TempFile() as tmp_file:
+        pickle_obj(obj=obj, filepath=tmp_file.name, serializer=serializer)
+        obj_from_disk = unpickle_obj(filepath=tmp_file.name, serializer=serializer)
+    return obj_from_disk  # noqa: RET504
+
+
+def assert_silent(capsys) -> None:
+    """
+    Given a ``CaptureFixture`` instance (from the ``pytest`` built-in ``capsys`` fixture),
+    read the recently-captured data into a variable and assert that nothing was written
+    to stdout or stderr.
+
+    This is just here to turn 3 lines of repetitive code into 1.
+
+    Note that this does have a side effect... ``capsys.readouterr()`` copies
+    from a buffer then frees it. So it will only store into ``.out`` and ``.err`` the
+    captured output since the last time that ``.readouterr()`` was called.
+
+    ref: https://docs.pytest.org/en/stable/how-to/capture-stdout-stderr.html
+    """
+    captured = capsys.readouterr()
+    assert captured.out == "", captured.out
+    assert captured.err == "", captured.err
+
+
+# doing this here, at import time, to ensure it only runs once_per import
+# instead of once per assertion
+_numpy_testing_supports_strict_kwarg = "strict" in getfullargspec(np.testing.assert_array_equal).kwonlyargs
+
+
+def np_assert_array_equal(*args, **kwargs):
+    """
+    np.testing.assert_array_equal() only got the kwarg ``strict`` in June 2022:
+    https://github.com/numpy/numpy/pull/21595
+
+    This function is here for testing on older Python (and therefore ``numpy``)
+    """
+    if not _numpy_testing_supports_strict_kwarg:
+        kwargs.pop("strict")
+    np.testing.assert_array_equal(*args, **kwargs)
+
+
+def assert_subtree_valid(root):
+    """Recursively checks the validity of a subtree rooted at `root`.
+
+    Currently it only checks whether weights and counts are consistent between
+    all parent nodes and their children.
+
+    Parameters
+    ----------
+    root : dict
+        A dictionary representing the root of the subtree.
+        It should be produced by dump_model()
+
+    Returns
+    -------
+    tuple
+        A tuple containing the weight and count of the subtree rooted at `root`.
+    """
+    if "leaf_count" in root:
+        return (root["leaf_weight"], root["leaf_count"])
+
+    left_child = root["left_child"]
+    right_child = root["right_child"]
+    (l_w, l_c) = assert_subtree_valid(left_child)
+    (r_w, r_c) = assert_subtree_valid(right_child)
+    assert abs(root["internal_weight"] - (l_w + r_w)) <= 1e-3, (
+        "root node's internal weight should be approximately the sum of its child nodes' internal weights"
+    )
+    assert root["internal_count"] == l_c + r_c, (
+        "root node's internal count should be exactly the sum of its child nodes' internal counts"
+    )
+    return (root["internal_weight"], root["internal_count"])
+
+
+def assert_all_trees_valid(model_dump):
+    for idx, tree in enumerate(model_dump["tree_info"]):
+        assert tree["tree_index"] == idx, f"tree {idx} should have tree_index={idx}. Full tree: {tree}"
+        assert_subtree_valid(tree["tree_structure"])
+
+
+# This mapping from CI-time environment variables is a placeholder
+# until there is a more reliable way to detect which customizations
+# LightGBM was built with.
+#
+# see https://github.com/lightgbm-org/LightGBM/issues/7273
+#
+class BuildInfo:
+    has_cuda = os.getenv("TASK", "") == "cuda"
+    has_gpu = os.getenv("TASK", "") == "gpu"
+    has_mpi = os.getenv("TASK", "") == "mpi"