[c++][python-package] add int8 input dtype for pre-discretized features

include/LightGBM/c_api.h

@@ -37,6 +37,7 @@ typedef void* ByteBufferHandle; /*!< \brief Handle of ByteBuffer. */
 #define C_API_DTYPE_FLOAT64 (1)  /*!< \brief float64 (double precision float). */
 #define C_API_DTYPE_INT32   (2)  /*!< \brief int32. */
 #define C_API_DTYPE_INT64   (3)  /*!< \brief int64. */
+#define C_API_DTYPE_INT8    (4)  /*!< \brief int8. */
 
 #define C_API_PREDICT_NORMAL     (0)  /*!< \brief Normal prediction, with transform (if needed). */
 #define C_API_PREDICT_RAW_SCORE  (1)  /*!< \brief Predict raw score. */

python-package/lightgbm/basic.py

@@ -77,6 +77,7 @@
 _ctypes_float_ptr = Union[
     "ctypes._Pointer[ctypes.c_float]",
     "ctypes._Pointer[ctypes.c_double]",
+    "ctypes._Pointer[ctypes.c_int8]",
 ]
 _ctypes_float_array = Union[
     "ctypes.Array[ctypes._Pointer[ctypes.c_float]]",
@@ -197,7 +198,7 @@ def _get_sample_count(total_nrow: int, params: str) -> int:
 
 def _np2d_to_np1d(mat: np.ndarray) -> Tuple[np.ndarray, int]:
     dtype: "np.typing.DTypeLike"
-    if mat.dtype in (np.float32, np.float64):
+    if mat.dtype in (np.float32, np.float64, np.int8):
         dtype = mat.dtype
     else:
         dtype = np.float32
@@ -701,6 +702,7 @@ def _choose_param_value(main_param_name: str, params: Dict[str, Any], default_va
 _C_API_DTYPE_FLOAT64 = 1
 _C_API_DTYPE_INT32 = 2
 _C_API_DTYPE_INT64 = 3
+_C_API_DTYPE_INT8 = 4
 
 """Macro definition of data order in matrix"""
 _C_API_IS_COL_MAJOR = 0
@@ -749,7 +751,7 @@ def _convert_from_sliced_object(data: np.ndarray) -> np.ndarray:
 
 
 def _c_float_array(data: np.ndarray) -> Tuple[_ctypes_float_ptr, int, np.ndarray]:
-    """Get pointer of float numpy array / list."""
+    """Get pointer of float/int8 numpy array / list."""
     if _is_1d_list(data):
         data = np.asarray(data)
     if _is_numpy_1d_array(data):
@@ -762,8 +764,11 @@ def _c_float_array(data: np.ndarray) -> Tuple[_ctypes_float_ptr, int, np.ndarray
         elif data.dtype == np.float64:
             ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
             type_data = _C_API_DTYPE_FLOAT64
+        elif data.dtype == np.int8:
+            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int8))
+            type_data = _C_API_DTYPE_INT8
         else:
-            raise TypeError(f"Expected np.float32 or np.float64, met type({data.dtype})")
+            raise TypeError(f"Expected np.float32, np.float64 or np.int8, met type({data.dtype})")
     else:
         raise TypeError(f"Unknown type({type(data).__name__})")
     return (ptr_data, type_data, data)  # return `data` to avoid the temporary copy is freed

src/c_api.cpp

@@ -2856,6 +2856,8 @@ RowFunctionFromDenseMatrix(const void* data, int num_row, int num_col, int data_
     return RowFunctionFromDenseMatrix_helper<float>(data, num_row, num_col, is_row_major);
   } else if (data_type == C_API_DTYPE_FLOAT64) {
     return RowFunctionFromDenseMatrix_helper<double>(data, num_row, num_col, is_row_major);
+  } else if (data_type == C_API_DTYPE_INT8) {
+    return RowFunctionFromDenseMatrix_helper<int8_t>(data, num_row, num_col, is_row_major);
   }
   Log::Fatal("Unknown data type in RowFunctionFromDenseMatrix");
   return nullptr;

tests/python_package_test/test_basic.py

@@ -1020,7 +1020,7 @@ def test_max_depth_warning_is_raised_if_max_depth_gte_5_and_num_leaves_omitted(c
 
 
 @pytest.mark.parametrize("order", ["C", "F"])
-@pytest.mark.parametrize("dtype", ["float32", "int64"])
+@pytest.mark.parametrize("dtype", ["float32", "int8", "int64"])
 def test_no_copy_in_dataset_from_numpy_2d(rng, order, dtype):
     X = rng.random(size=(100, 3))
     X = np.require(X, dtype=dtype, requirements=order)
@@ -1029,13 +1029,44 @@ def test_no_copy_in_dataset_from_numpy_2d(rng, order, dtype):
         assert layout == lgb.basic._C_API_IS_COL_MAJOR
     else:
         assert layout == lgb.basic._C_API_IS_ROW_MAJOR
-    if dtype == "float32":
+    if dtype in ("float32", "int8"):
         assert np.shares_memory(X, X1d)
     else:
         # makes a copy
         assert not np.shares_memory(X, X1d)
 
 
+def test_c_float_array_accepts_int8():
+    data = np.array([0, 1, 2, 3, 4], dtype=np.int8)
+    _, type_data, holder = lgb.basic._c_float_array(data)
+    assert type_data == lgb.basic._C_API_DTYPE_INT8
+    # int8 input is forwarded without a float conversion
+    assert holder.dtype == np.int8
+    assert np.shares_memory(data, holder)
+
+
+@pytest.mark.parametrize("dtype", [np.int16, np.int32, np.int64, np.uint8])
+def test_c_float_array_rejects_unsupported_int_dtypes(dtype):
+    data = np.array([0, 1, 2], dtype=dtype)
+    with pytest.raises(TypeError, match=r"Expected np\.float32, np\.float64 or np\.int8"):
+        lgb.basic._c_float_array(data)
+
+
+def test_dataset_from_int8_matches_float32(rng, tmp_path):
+    # values fit losslessly in both int8 and float32, so the resulting
+    # binned datasets must be identical regardless of the input dtype
+    X_int8 = rng.integers(low=0, high=5, size=(500, 4)).astype(np.int8)
+    y = rng.integers(low=0, high=2, size=500).astype(np.float64)
+
+    int8_path = tmp_path / "int8.txt"
+    lgb.Dataset(X_int8, y)._dump_text(int8_path)
+
+    float32_path = tmp_path / "float32.txt"
+    lgb.Dataset(X_int8.astype(np.float32), y)._dump_text(float32_path)
+
+    assert filecmp.cmp(int8_path, float32_path)
+
+
 def test_equal_datasets_from_row_major_and_col_major_data(tmp_path):
     # row-major dataset
     X_row, y = make_blobs(n_samples=1_000, n_features=3, centers=2)