java-topology/defects/sklearn/patch/sklearn-0001-gradient-boosting-feature-names-index.md

UNDF: UNDF-2026-000000715

sklearn-0001: HistGradientBoosting _check_categories — O(n²) feature_names.index in loop

Severity: MEDIUM CWE: CWE-407 (Algorithmic Complexity — linear membership test in hot loop) Speedup: >40x at F=1000 features, C=200 categorical features Target: scikit-learn (scikit-learn/scikit-learn) File: sklearn/ensemble/_hist_gradient_boosting/gradient_boosting.py:440-443

Description

_check_categories resolves categorical feature names to integer indices by calling feature_names.index(feature_name) inside a for feature_name in categorical_features loop. feature_names is a plain Python list built from the DataFrame column names. Each call to .index() performs a linear scan from position 0, making the total complexity O(C × F) where C is the number of categorical features and F is the total feature count.

With wide feature-rich datasets (1000+ columns, 200+ categorical), this is called at every fit(), predict(), and score() invocation — directly on the hot path for model training and inference.

Root Cause

# sklearn/ensemble/_hist_gradient_boosting/gradient_boosting.py:440-443
is_categorical = np.zeros(n_features, dtype=bool)
feature_names = list(feature_names_in_)          # plain list
for feature_name in categorical_features:
    try:
        is_categorical[feature_names.index(feature_name)] = True  # O(F) per cat feature
    except ValueError as e:
        ...

Fix: build a dict mapping name → index once before the loop.

Patch

is_categorical = np.zeros(n_features, dtype=bool)
feature_names = list(feature_names_in_)
feature_name_to_idx = {name: i for i, name in enumerate(feature_names)}  # O(F) once
for feature_name in categorical_features:
    try:
        is_categorical[feature_name_to_idx[feature_name]] = True           # O(1)
    except KeyError:
        raise ValueError(
            f"categorical_features has an item value '{feature_name}' "
            "which is not a valid feature name of the training "
            f"data. Observed feature names: {feature_names}"
        )

Complexity Before

O(C × F) — C categorical features × F total features per index lookup

Complexity After

Build index once: O(F), then O(1) per lookup → O(F + C) total

Reproduction

cd defects/sklearn/unit && javac -d . *.java && java -ea unit.SklearnTest