From edca9b931b49e8306d978fec21ab98deb17f32b9 Mon Sep 17 00:00:00 2001 From: Max Zwager Date: Wed, 24 Jun 2026 13:53:45 +0200 Subject: [PATCH] make Dataset.set_position update self.position consistently with set_weight --- python-package/lightgbm/basic.py | 1 + tests/python_package_test/test_basic.py | 44 ++++++++++++++++++++++-- tests/python_package_test/test_engine.py | 7 ++-- 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py index d32cec47eef0..258ed0158d41 100644 --- a/python-package/lightgbm/basic.py +++ b/python-package/lightgbm/basic.py @@ -3114,6 +3114,7 @@ def set_position( if self._handle is not None and position is not None: position = _list_to_1d_numpy(data=position, dtype=np.int32, name="position") self.set_field("position", position) + self.position = self.get_field("position") # original values can be modified at cpp side return self def get_feature_name(self) -> List[str]: diff --git a/tests/python_package_test/test_basic.py b/tests/python_package_test/test_basic.py index c66b280041f4..7313bb196a92 100644 --- a/tests/python_package_test/test_basic.py +++ b/tests/python_package_test/test_basic.py @@ -605,11 +605,13 @@ def test_dataset_construction_overwrites_user_provided_metadata_fields(): np_assert_array_equal(dtrain.get_field("label"), expected_label, strict=True) if not BuildInfo.has_cuda: - expected_position = np.array([0.0, 1.0], dtype=np.float32) + # NOTE: "position" is converted to int32 on the C++ side and remapped to dense + # internal indices in encounter order. Here the input [0, 1] is already dense + # starting from 0 in encounter order, so the remap is the identity. + expected_position = np.array([0, 1], dtype=np.int32) np_assert_array_equal(dtrain.position, expected_position, strict=True) np_assert_array_equal(dtrain.get_position(), expected_position, strict=True) - # NOTE: "position" is converted to int32 on the C++ side - np_assert_array_equal(dtrain.get_field("position"), np.array([0.0, 1.0], dtype=np.int32), strict=True) + np_assert_array_equal(dtrain.get_field("position"), expected_position, strict=True) expected_weight = np.array([0.5, 1.5], dtype=np.float32) np_assert_array_equal(dtrain.weight, expected_weight, strict=True) @@ -617,6 +619,42 @@ def test_dataset_construction_overwrites_user_provided_metadata_fields(): np_assert_array_equal(dtrain.get_field("weight"), expected_weight, strict=True) +@pytest.mark.skipif( + BuildInfo.has_cuda, + reason="Positions in learning to rank is not supported in CUDA version yet", +) +def test_set_position_updates_self_position_with_remapped_int32_values(): + # Position values are remapped to dense int32 indices in the order they are first + # encountered. With input [3, 1, 0, 2, 4, 3, 1, 0, 2, 4]: + # 3 -> 0 (first encountered), 1 -> 1, 0 -> 2, 2 -> 3, 4 -> 4 + X = np.arange(20, dtype=np.float64).reshape(10, 2) + y = np.arange(10, dtype=np.float64) + position = np.array([3, 1, 0, 2, 4, 3, 1, 0, 2, 4], dtype=np.int64) + expected = np.array([0, 1, 2, 3, 4, 0, 1, 2, 3, 4], dtype=np.int32) + + # set via constructor + dtrain = lgb.Dataset( + X, + label=y, + position=position, + params={"min_data_in_bin": 1, "min_data_in_leaf": 1, "verbosity": -1}, + ).construct() + np_assert_array_equal(dtrain.position, expected, strict=True) + np_assert_array_equal(dtrain.get_position(), expected, strict=True) + np_assert_array_equal(dtrain.get_field("position"), expected, strict=True) + + # set via set_position() on an already-constructed Dataset + dtrain2 = lgb.Dataset( + X, + label=y, + params={"min_data_in_bin": 1, "min_data_in_leaf": 1, "verbosity": -1}, + ).construct() + dtrain2.set_position(position) + np_assert_array_equal(dtrain2.position, expected, strict=True) + np_assert_array_equal(dtrain2.get_position(), expected, strict=True) + np_assert_array_equal(dtrain2.get_field("position"), expected, strict=True) + + def test_dataset_construction_with_high_cardinality_categorical_succeeds(rng): pd = pytest.importorskip("pandas") X = pd.DataFrame({"x1": rng.integers(low=0, high=5_000, size=(10_000,))}) diff --git a/tests/python_package_test/test_engine.py b/tests/python_package_test/test_engine.py index 5331361a22f4..75c24d190f5e 100644 --- a/tests/python_package_test/test_engine.py +++ b/tests/python_package_test/test_engine.py @@ -848,9 +848,12 @@ def test_ranking_with_position_information_with_dataset_constructor(tmp_path): gbm_unbiased_set_position = lgb.train(params, lgb_train, valid_sets=lgb_valid, num_boost_round=50) assert gbm_unbiased.best_score["valid_0"]["ndcg@3"] == gbm_unbiased_set_position.best_score["valid_0"]["ndcg@3"] - # test get_position works + # test get_position works (positions are remapped to dense int32 indices on the C++ + # side, so compare against get_field("position") rather than the original input) positions_from_get = lgb_train.get_position() - np_assert_array_equal(positions_from_get, positions, strict=True) + np_assert_array_equal(positions_from_get, lgb_train.get_field("position"), strict=True) + assert positions_from_get.dtype == np.int32 + assert positions_from_get.shape == positions.shape def test_early_stopping():