[python-package] Allow to pass Arrow array & Polars series as position

include/LightGBM/c_api.h

@@ -563,11 +563,11 @@ LIGHTGBM_C_EXPORT int LGBM_DatasetDumpText(DatasetHandle handle,
 /*!
  * \brief Set vector to a content in info.
  * \note
- * - \a group only works for ``C_API_DTYPE_INT32``;
+ * - \a group and \a position only work for ``C_API_DTYPE_INT32``;
  * - \a label and \a weight only work for ``C_API_DTYPE_FLOAT32``;
  * - \a init_score only works for ``C_API_DTYPE_FLOAT64``.
  * \param handle Handle of dataset
- * \param field_name Field name, can be \a label, \a weight, \a init_score, \a group
+ * \param field_name Field name, can be \a label, \a weight, \a init_score, \a group, \a position
  * \param field_data Pointer to data vector
  * \param num_element Number of elements in ``field_data``
  * \param type Type of ``field_data`` pointer, can be ``C_API_DTYPE_INT32``, ``C_API_DTYPE_FLOAT32`` or ``C_API_DTYPE_FLOAT64``
@@ -583,11 +583,11 @@ LIGHTGBM_C_EXPORT int LGBM_DatasetSetField(DatasetHandle handle,
  * \brief Set vector to a content in info.
  * \deprecated This function is deprecated in favor of ``LGBM_DatasetSetFieldFromArrowStream``.
  * \note
- * - \a group converts input datatype into ``int32``;
+ * - \a group and \a position convert input datatype into ``int32``;
  * - \a label and \a weight convert input datatype into ``float32``;
  * - \a init_score converts input datatype into ``float64``.
  * \param handle Handle of dataset
- * \param field_name Field name, can be \a label, \a weight, \a init_score, \a group
+ * \param field_name Field name, can be \a label, \a weight, \a init_score, \a group, \a position
  * \param n_chunks The number of Arrow arrays passed to this function
  * \param chunks Pointer to the list of Arrow arrays
  * \param schema Pointer to the schema of all Arrow arrays

include/LightGBM/dataset.h

@@ -124,6 +124,8 @@ class Metadata {
   void SetQuery(int64_t n_chunks, struct ArrowArray* chunks, struct ArrowSchema* schema);
 
   void SetPosition(const data_size_t* position, data_size_t len);
+  void SetPosition(struct ArrowArrayStream* stream);
+  void SetPosition(int64_t n_chunks, struct ArrowArray* chunks, struct ArrowSchema* schema);
 
   /*!
   * \brief Set initial scores

python-package/lightgbm/basic.py

@@ -75,9 +75,12 @@
     pd_Series,
     nwt.IntoSeries,
 ]
+# 'position' intentionally does not support 'list' inputs
+# ref: https://github.com/lightgbm-org/LightGBM/pull/5929#discussion_r1262646998
 _LGBM_PositionType = Union[
     np.ndarray,
     pd_Series,
+    nwt.IntoSeries,
 ]
 _LGBM_InitScoreType = Union[
     List[float],
@@ -1752,7 +1755,7 @@ def __init__(
             Other parameters for Dataset.
         free_raw_data : bool, optional (default=True)
             If True, raw data is freed after constructing inner Dataset.
-        position : numpy 1-D array, pandas Series or None, optional (default=None)
+        position : numpy 1-D array, pandas Series, pyarrow ChunkedArray, polars Series or None, optional (default=None)
             Position of items used in unbiased learning-to-rank task.
         """
         self._handle: Optional[_DatasetHandle] = None
@@ -2577,7 +2580,7 @@ def create_valid(
             Init score for Dataset.
         params : dict or None, optional (default=None)
             Other parameters for validation Dataset.
-        position : numpy 1-D array, pandas Series or None, optional (default=None)
+        position : numpy 1-D array, pandas Series, pyarrow ChunkedArray, polars Series or None, optional (default=None)
             Position of items used in unbiased learning-to-rank task.
 
         Returns
@@ -3102,7 +3105,7 @@ def set_position(
 
         Parameters
         ----------
-        position : numpy 1-D array, pandas Series or None, optional (default=None)
+        position : numpy 1-D array, pandas Series, pyarrow ChunkedArray, polars Series or None, optional (default=None)
             Position of items used in unbiased learning-to-rank task.
 
         Returns
@@ -3112,7 +3115,8 @@ def set_position(
         """
         self.position = position
         if self._handle is not None and position is not None:
-            position = _list_to_1d_numpy(data=position, dtype=np.int32, name="position")
+            if isinstance(position, pd_Series) or not nwd.is_into_series(position):
+                position = _list_to_1d_numpy(data=position, dtype=np.int32, name="position")
             self.set_field("position", position)
         return self
 
@@ -3261,7 +3265,7 @@ def get_position(self) -> Optional[_LGBM_PositionType]:
 
         Returns
         -------
-        position : numpy 1-D array, pandas Series or None
+        position : numpy 1-D array, pandas Series, pyarrow ChunkedArray, polars Series or None
             Position of items used in unbiased learning-to-rank task.
             For a constructed ``Dataset``, this will only return ``None`` or a numpy array.
         """

src/io/dataset.cpp

@@ -941,6 +941,8 @@ bool Dataset::SetFieldFromArrow(const char* field_name, struct ArrowArrayStream*
     metadata_.SetInitScore(stream);
   } else if (name == std::string("query") || name == std::string("group")) {
     metadata_.SetQuery(stream);
+  } else if (name == std::string("position")) {
+    metadata_.SetPosition(stream);
   } else {
     return false;
   }
@@ -959,6 +961,8 @@ bool Dataset::SetFieldFromArrow(const char* field_name, int64_t n_chunks,
     metadata_.SetInitScore(n_chunks, chunks, schema);
   } else if (name == std::string("query") || name == std::string("group")) {
     metadata_.SetQuery(n_chunks, chunks, schema);
+  } else if (name == std::string("position")) {
+    metadata_.SetPosition(n_chunks, chunks, schema);
   } else {
     return false;
   }

src/io/metadata.cpp

@@ -691,6 +691,33 @@ void Metadata::SetPosition(const data_size_t* positions, data_size_t len) {
   }
 }
 
+#ifndef LGB_R_BUILD
+void Metadata::SetPosition(struct ArrowArrayStream* stream) {
+  ArrowChunkedArray chunked_array(stream);
+  chunked_array.view().visit<data_size_t>([&](auto&& visitor) {
+    std::vector<data_size_t> positions;
+    positions.reserve(visitor.end() - visitor.begin());
+    for (auto it = visitor.begin(); it != visitor.end(); ++it) {
+      positions.push_back(*it);
+    }
+    SetPosition(positions.data(), static_cast<data_size_t>(positions.size()));
+  });
+}
+
+void Metadata::SetPosition(int64_t n_chunks, struct ArrowArray* chunks,
+                           struct ArrowSchema* schema) {
+  ArrowChunkedArray chunked_array(n_chunks, chunks, schema);
+  chunked_array.view().visit<data_size_t>([&](auto&& visitor) {
+    std::vector<data_size_t> positions;
+    positions.reserve(visitor.end() - visitor.begin());
+    for (auto it = visitor.begin(); it != visitor.end(); ++it) {
+      positions.push_back(*it);
+    }
+    SetPosition(positions.data(), static_cast<data_size_t>(positions.size()));
+  });
+}
+#endif  // LGB_R_BUILD
+
 void Metadata::InsertQueries(const data_size_t* queries, data_size_t start_index, data_size_t len) {
   if (!queries) {
     Log::Fatal("Passed null queries");

tests/python_package_test/test_arrow.py

@@ -293,6 +293,42 @@ def test_dataset_construct_groups(group_data, arrow_type):
     np_assert_array_equal(expected, dataset.get_field("group"), strict=True)
 
 
+# ------------------------------------------ POSITION ------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "position_data",
+    [
+        [[0, 1, 2, 3, 4]],
+        [[0, 1, 2], [3, 4]],
+        [[], [0, 1, 2], [3, 4]],
+        [[0, 1], [], [2], [3, 4], []],
+    ],
+)
+@pytest.mark.parametrize("arrow_type", _INTEGER_TYPES)
+def test_dataset_construct_position(position_data, arrow_type):
+    data = generate_dummy_arrow_table()
+    positions = pa.chunked_array(position_data, type=arrow_type)
+    dataset = lgb.Dataset(data, label=[0, 1, 0, 1, 0], position=positions, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 2, 3, 4], dtype=np.int32)
+    np_assert_array_equal(expected, dataset.get_field("position"), strict=True)
+
+
+@pytest.mark.parametrize("arrow_type", _INTEGER_TYPES)
+def test_dataset_construct_position_with_duplicates_and_out_of_order(arrow_type):
+    data = generate_dummy_arrow_table()
+    positions = pa.chunked_array([[15, 15, 8, 27, 15]], type=arrow_type)
+    dataset = lgb.Dataset(data, label=[0, 1, 0, 1, 0], position=positions, params=dummy_dataset_params())
+    dataset.construct()
+
+    # positions are remapped on the C++ side to dense indices in first-seen order:
+    # 15 -> 0, 8 -> 1, 27 -> 2
+    expected = np.array([0, 0, 1, 2, 0], dtype=np.int32)
+    np_assert_array_equal(expected, dataset.get_field("position"), strict=True)
+
+
 # ----------------------------------------- INIT SCORES ----------------------------------------- #
 
 

tests/python_package_test/test_polars.py

@@ -229,6 +229,33 @@ def test_dataset_construct_groups(polars_type):
     np_assert_array_equal(expected, dataset.get_field("group"), strict=True)
 
 
+# ------------------------------------------ POSITION ------------------------------------------- #
+
+
+@pytest.mark.parametrize("polars_type", _INTEGER_TYPES)
+def test_dataset_construct_position(polars_type):
+    data = generate_dummy_polars_frame()
+    positions = pl.Series("position", [0, 1, 2, 3, 4], dtype=polars_type)
+    dataset = lgb.Dataset(data, label=[0, 1, 0, 1, 0], position=positions, params=dummy_dataset_params())
+    dataset.construct()
+
+    expected = np.array([0, 1, 2, 3, 4], dtype=np.int32)
+    np_assert_array_equal(expected, dataset.get_field("position"), strict=True)
+
+
+@pytest.mark.parametrize("polars_type", _INTEGER_TYPES)
+def test_dataset_construct_position_with_duplicates_and_out_of_order(polars_type):
+    data = generate_dummy_polars_frame()
+    positions = pl.Series("position", [15, 15, 8, 27, 15], dtype=polars_type)
+    dataset = lgb.Dataset(data, label=[0, 1, 0, 1, 0], position=positions, params=dummy_dataset_params())
+    dataset.construct()
+
+    # positions are remapped on the C++ side to dense indices in first-seen order:
+    # 15 -> 0, 8 -> 1, 27 -> 2
+    expected = np.array([0, 0, 1, 2, 0], dtype=np.int32)
+    np_assert_array_equal(expected, dataset.get_field("position"), strict=True)
+
+
 # ----------------------------------------- INIT SCORES ----------------------------------------- #