MNT move bitset from HGBT to utils (scikit-learn#33957)

lorentzenchr · adam2392 · web-flow · commit e02dfc40d632 · 2026-05-09T10:05:18.000-04:00
Co-authored-by: Adam Li &lt;adam2392@gmail.com&gt;
diff --git a/sklearn/ensemble/_hist_gradient_boosting/_bitset.pxd b/sklearn/ensemble/_hist_gradient_boosting/_bitset.pxd
diff --git a/sklearn/ensemble/_hist_gradient_boosting/_bitset.pyx b/sklearn/ensemble/_hist_gradient_boosting/_bitset.pyx
diff --git a/sklearn/ensemble/_hist_gradient_boosting/_predictor.pyx b/sklearn/ensemble/_hist_gradient_boosting/_predictor.pyx
@@ -5,14 +5,13 @@ from cython.parallel import prange
 from libc.math cimport isnan
 import numpy as np
 
+from sklearn.utils._bitset cimport BITSET_INNER_DTYPE_C, in_bitset_2d_memoryview
 from sklearn.utils._typedefs cimport intp_t, uint8_t
 from sklearn.ensemble._hist_gradient_boosting.common cimport X_DTYPE_C
 from sklearn.ensemble._hist_gradient_boosting.common cimport Y_DTYPE_C
 from sklearn.ensemble._hist_gradient_boosting.common import Y_DTYPE
 from sklearn.ensemble._hist_gradient_boosting.common cimport X_BINNED_DTYPE_C
-from sklearn.ensemble._hist_gradient_boosting.common cimport BITSET_INNER_DTYPE_C
 from sklearn.ensemble._hist_gradient_boosting.common cimport node_struct
-from sklearn.ensemble._hist_gradient_boosting._bitset cimport in_bitset_2d_memoryview
 
 
 def _predict_from_raw_data(  # raw data = non-binned data
diff --git a/sklearn/ensemble/_hist_gradient_boosting/binning.py b/sklearn/ensemble/_hist_gradient_boosting/binning.py
@@ -14,14 +14,14 @@
 
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.ensemble._hist_gradient_boosting._binning import _map_to_bins
-from sklearn.ensemble._hist_gradient_boosting._bitset import set_bitset_memoryview
 from sklearn.ensemble._hist_gradient_boosting.common import (
     ALMOST_INF,
     X_BINNED_DTYPE,
     X_BITSET_INNER_DTYPE,
     X_DTYPE,
 )
 from sklearn.utils import check_array, check_random_state
+from sklearn.utils._bitset import set_bitset_memoryview
 from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
 from sklearn.utils.parallel import Parallel, delayed
 from sklearn.utils.stats import _weighted_percentile
diff --git a/sklearn/ensemble/_hist_gradient_boosting/common.pxd b/sklearn/ensemble/_hist_gradient_boosting/common.pxd
@@ -1,12 +1,14 @@
-from sklearn.utils._typedefs cimport float32_t, float64_t, intp_t, uint8_t, uint32_t
+# Authors: The scikit-learn developers
+# SPDX-License-Identifier: BSD-3-Clause
+
+from sklearn.utils._bitset cimport BITSET_INNER_DTYPE_C, BITSET_INNER_DTYPE_C
+from sklearn.utils._typedefs cimport float32_t, float64_t, intp_t, uint8_t
 
 
 ctypedef float64_t X_DTYPE_C
 ctypedef uint8_t X_BINNED_DTYPE_C
 ctypedef float64_t Y_DTYPE_C
 ctypedef float32_t G_H_DTYPE_C
-ctypedef uint32_t BITSET_INNER_DTYPE_C
-ctypedef BITSET_INNER_DTYPE_C[8] BITSET_DTYPE_C
 
 
 cdef packed struct hist_struct:
diff --git a/sklearn/ensemble/_hist_gradient_boosting/common.pyx b/sklearn/ensemble/_hist_gradient_boosting/common.pyx
@@ -1,3 +1,6 @@
+# Authors: The scikit-learn developers
+# SPDX-License-Identifier: BSD-3-Clause
+
 import numpy as np
 
 # Y_DYTPE is the dtype to which the targets y are converted to. This is also
diff --git a/sklearn/ensemble/_hist_gradient_boosting/grower.py b/sklearn/ensemble/_hist_gradient_boosting/grower.py
@@ -14,9 +14,6 @@
 
 import numpy as np
 
-from sklearn.ensemble._hist_gradient_boosting._bitset import (
-    set_raw_bitset_from_binned_bitset,
-)
 from sklearn.ensemble._hist_gradient_boosting.common import (
     PREDICTOR_RECORD_DTYPE,
     X_BITSET_INNER_DTYPE,
@@ -25,6 +22,7 @@
 from sklearn.ensemble._hist_gradient_boosting.histogram import HistogramBuilder
 from sklearn.ensemble._hist_gradient_boosting.predictor import TreePredictor
 from sklearn.ensemble._hist_gradient_boosting.splitting import Splitter
+from sklearn.utils._bitset import set_raw_bitset_from_binned_bitset
 from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
 
 
diff --git a/sklearn/ensemble/_hist_gradient_boosting/meson.build b/sklearn/ensemble/_hist_gradient_boosting/meson.build
@@ -12,7 +12,6 @@ hist_gradient_boosting_extension_metadata = {
   'splitting': {'sources': [cython_gen.process('splitting.pyx')], 'dependencies': [openmp_dep]},
   '_binning': {'sources': [cython_gen.process('_binning.pyx')], 'dependencies': [openmp_dep]},
   '_predictor': {'sources': [cython_gen.process('_predictor.pyx')], 'dependencies': [openmp_dep]},
-  '_bitset': {'sources': [cython_gen.process('_bitset.pyx')]},
   'common': {'sources': [cython_gen.process('common.pyx')]},
 }
 
diff --git a/sklearn/ensemble/_hist_gradient_boosting/splitting.pyx b/sklearn/ensemble/_hist_gradient_boosting/splitting.pyx
@@ -16,16 +16,13 @@ from libc.math cimport INFINITY, ceil
 from libc.stdlib cimport malloc, free, qsort
 from libc.string cimport memcpy
 
+from sklearn.utils._bitset cimport BITSET_DTYPE_C, BITSET_INNER_DTYPE_C
+from sklearn.utils._bitset cimport in_bitset, init_bitset, set_bitset
 from sklearn.utils._typedefs cimport uint8_t
 from sklearn.ensemble._hist_gradient_boosting.common cimport X_BINNED_DTYPE_C
 from sklearn.ensemble._hist_gradient_boosting.common cimport Y_DTYPE_C
 from sklearn.ensemble._hist_gradient_boosting.common cimport hist_struct
-from sklearn.ensemble._hist_gradient_boosting.common cimport BITSET_INNER_DTYPE_C
-from sklearn.ensemble._hist_gradient_boosting.common cimport BITSET_DTYPE_C
 from sklearn.ensemble._hist_gradient_boosting.common cimport MonotonicConstraint
-from sklearn.ensemble._hist_gradient_boosting._bitset cimport init_bitset
-from sklearn.ensemble._hist_gradient_boosting._bitset cimport set_bitset
-from sklearn.ensemble._hist_gradient_boosting._bitset cimport in_bitset
 
 
 cdef struct split_info_struct:
diff --git a/sklearn/ensemble/_hist_gradient_boosting/tests/test_predictor.py b/sklearn/ensemble/_hist_gradient_boosting/tests/test_predictor.py
@@ -3,10 +3,6 @@
 from numpy.testing import assert_allclose
 
 from sklearn.datasets import make_regression
-from sklearn.ensemble._hist_gradient_boosting._bitset import (
-    set_bitset_memoryview,
-    set_raw_bitset_from_binned_bitset,
-)
 from sklearn.ensemble._hist_gradient_boosting.binning import _BinMapper
 from sklearn.ensemble._hist_gradient_boosting.common import (
     ALMOST_INF,
@@ -20,6 +16,10 @@
 from sklearn.ensemble._hist_gradient_boosting.predictor import TreePredictor
 from sklearn.metrics import r2_score
 from sklearn.model_selection import train_test_split
+from sklearn.utils._bitset import (
+    set_bitset_memoryview,
+    set_raw_bitset_from_binned_bitset,
+)
 from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
 
 n_threads = _openmp_effective_n_threads()
diff --git a/sklearn/utils/_bitset.pxd b/sklearn/utils/_bitset.pxd
@@ -0,0 +1,21 @@
+# Authors: The scikit-learn developers
+# SPDX-License-Identifier: BSD-3-Clause
+
+from sklearn.utils._typedefs cimport float64_t, uint8_t, uint32_t
+
+ctypedef uint32_t BITSET_INNER_DTYPE_C
+ctypedef BITSET_INNER_DTYPE_C[8] BITSET_DTYPE_C
+
+cdef void init_bitset(BITSET_DTYPE_C bitset) noexcept nogil
+
+cdef void set_bitset(BITSET_DTYPE_C bitset, uint8_t val) noexcept nogil
+
+cdef uint8_t in_bitset(BITSET_DTYPE_C bitset, uint8_t val) noexcept nogil
+
+cpdef uint8_t in_bitset_memoryview(
+    const BITSET_INNER_DTYPE_C[:] bitset, uint8_t val
+) noexcept nogil
+
+cdef uint8_t in_bitset_2d_memoryview(
+    const BITSET_INNER_DTYPE_C[:, :] bitset, uint8_t val, unsigned int row
+) noexcept nogil
diff --git a/sklearn/utils/_bitset.pyx b/sklearn/utils/_bitset.pyx
@@ -0,0 +1,63 @@
+"""
+A bitset is a data structure used to represent sets of integers in [0, n]. For decision
+trees, we use them to represent sets of features indices (e.g. features that go to the
+left child, or features that are categorical). For familiarity with bitsets and bitwise
+operations see:
+https://en.wikipedia.org/wiki/Bit_array
+https://en.wikipedia.org/wiki/Bitwise_operation
+"""
+
+# Authors: The scikit-learn developers
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+cdef inline void init_bitset(BITSET_DTYPE_C bitset) noexcept nogil:
+    cdef:
+        unsigned int i
+
+    for i in range(8):
+        bitset[i] = 0
+
+
+cdef inline void set_bitset(BITSET_DTYPE_C bitset, uint8_t val) noexcept nogil:
+    bitset[val // 32] |= (1 << (val % 32))
+
+
+cdef inline uint8_t in_bitset(BITSET_DTYPE_C bitset, uint8_t val) noexcept nogil:
+    return (bitset[val // 32] >> (val % 32)) & 1
+
+
+cpdef inline uint8_t in_bitset_memoryview(
+    const BITSET_INNER_DTYPE_C[:] bitset, uint8_t val
+) noexcept nogil:
+    return (bitset[val // 32] >> (val % 32)) & 1
+
+
+cdef inline uint8_t in_bitset_2d_memoryview(
+    const BITSET_INNER_DTYPE_C[:, :] bitset, uint8_t val, unsigned int row
+) noexcept nogil:
+    # Same as above but works on 2d memory views to avoid the creation of 1d
+    # memory views. See https://github.com/scikit-learn/scikit-learn/issues/17299
+    return (bitset[row, val // 32] >> (val % 32)) & 1
+
+
+cpdef inline void set_bitset_memoryview(BITSET_INNER_DTYPE_C[:] bitset, uint8_t val):
+    bitset[val // 32] |= (1 << (val % 32))
+
+
+def set_raw_bitset_from_binned_bitset(
+    BITSET_INNER_DTYPE_C[:] raw_bitset,
+    BITSET_INNER_DTYPE_C[:] binned_bitset,
+    float64_t[:] categories,
+):
+    """Set the raw_bitset from the values of the binned bitset
+
+    categories is a mapping from binned category value to raw category value.
+    """
+    cdef:
+        int binned_cat_value
+        float64_t raw_cat_value
+
+    for binned_cat_value, raw_cat_value in enumerate(categories):
+        if in_bitset_memoryview(binned_bitset, binned_cat_value):
+            set_bitset_memoryview(raw_bitset, <uint8_t>raw_cat_value)
diff --git a/sklearn/utils/meson.build b/sklearn/utils/meson.build
@@ -24,6 +24,7 @@ utils_cython_tree = [
   # early in the build
   sklearn_root_cython_tree,
   fs.copyfile('__init__.py'),
+  fs.copyfile('_bitset.pxd'),
   fs.copyfile('_cython_blas.pxd'),
   fs.copyfile('_heap.pxd'),
   fs.copyfile('_openmp_helpers.pxd'),
@@ -38,6 +39,7 @@ utils_extension_metadata = {
     {'sources': [cython_gen.process('sparsefuncs_fast.pyx')]},
   '_cython_blas': {'sources': [cython_gen.process('_cython_blas.pyx'), _blas_int_pxi]},
   'arrayfuncs': {'sources': [cython_gen.process('arrayfuncs.pyx')]},
+  '_bitset': {'sources': [cython_gen.process('_bitset.pyx')]},
   'murmurhash': {
       'sources': [cython_gen.process('murmurhash.pyx'), 'src' / 'MurmurHash3.cpp'],
   },
diff --git a/sklearn/utils/tests/test_bitset.py b/sklearn/utils/tests/test_bitset.py
@@ -2,12 +2,11 @@
 import pytest
 from numpy.testing import assert_allclose
 
-from sklearn.ensemble._hist_gradient_boosting._bitset import (
+from sklearn.utils._bitset import (
     in_bitset_memoryview,
     set_bitset_memoryview,
     set_raw_bitset_from_binned_bitset,
 )
-from sklearn.ensemble._hist_gradient_boosting.common import X_DTYPE
 
 
 @pytest.mark.parametrize(
@@ -49,7 +48,7 @@ def test_raw_bitset_from_binned_bitset(
 ):
     binned_bitset = np.zeros(2, dtype=np.uint32)
     raw_bitset = np.zeros(2, dtype=np.uint32)
-    raw_categories = np.asarray(raw_categories, dtype=X_DTYPE)
+    raw_categories = np.asarray(raw_categories, dtype=np.float64)
 
     for val in binned_cat_to_insert:
         set_bitset_memoryview(binned_bitset, val)

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,6 @@ hist_gradient_boosting_extension_metadata = {`
`12`	`12`	`'splitting': {'sources': [cython_gen.process('splitting.pyx')], 'dependencies': [openmp_dep]},`
`13`	`13`	`'_binning': {'sources': [cython_gen.process('_binning.pyx')], 'dependencies': [openmp_dep]},`
`14`	`14`	`'_predictor': {'sources': [cython_gen.process('_predictor.pyx')], 'dependencies': [openmp_dep]},`
`15`		`- '_bitset': {'sources': [cython_gen.process('_bitset.pyx')]},`
`16`	`15`	`'common': {'sources': [cython_gen.process('common.pyx')]},`
`17`	`16`	`}`
`18`	`17`