Repository: scikit-learn/scikit-learn
File: sklearn/ensemble/_hist_gradient_boosting/binning.py (Python)

Code before:
    """
    # ignore missing values when computing bin thresholds
    missing_mask = np.isnan(col_data)
    if missing_mask.any():
        col_data = col_data[~missing_mask]

    # if sample weight is not None and null values exist
    # we need to remove those before calculating the
    # distinct points
    if sample_weight is not None:
        col_data_non_null = col_data[sample_weight != 0]
        distinct_values = np.unique(col_data_non_null).astype(X_DTYPE)
    else:
        distinct_values = np.unique(col_data).astype(X_DTYPE)

    # The data will be sorted anyway in np.unique and again in percentile, so we do it
    # here. Sorting also returns a contiguous array.
    sort_idx = np.argsort(col_data)
    col_data = col_data[sort_idx]

    # Calculate midpoints if distinct values <= max_bins
    if len(distinct_values) <= max_bins:
        midpoints = distinct_values[:-1] + distinct_values[1:]
        bin_thresholds = midpoints * 0.5

    elif sample_weight is None:
        percentiles = np.linspace(0, 100, num=max_bins + 1)
        percentiles = percentiles[1:-1]
        bin_thresholds = np.percentile(
            col_data, percentiles, method="averaged_inverted_cdf"
        )
        assert bin_thresholds.shape[0] == max_bins - 1

    else:
        # We could compute approximate midpoint percentiles using the output of
        # np.unique(col_data, return_counts) instead but this is more
        # work and the performance benefit will be limited because we
        # work on a fixed-size subsample of the full data.
        # TODO: check if there is a better way to implement this
        percentiles = np.linspace(0, 100, num=max_bins + 1)
        percentiles = percentiles[1:-1]
        sample_weight = sample_weight[~missing_mask]
        sample_weight = sample_weight[sort_idx]

        bin_thresholds = np.array(
            [
                _averaged_weighted_percentile(col_data, sample_weight, percentile)
                for percentile in percentiles
            ]
        )
        assert bin_thresholds.shape[0] == max_bins - 1

Review comment (performance):
Compared to the old code, it seems that you compute first the distinct values and then sort. But in the old code, sorting was done before, maybe for performance by giving already sorted arrays to np.unique and percentile. Did you loose some performance by sorting twice (in np.unique above and here) ?

EDIT: from some quick experiments, it doesn't seem to make much of a difference :)

Code after:
    Return
    ------
    binning_thresholds : ndarray of shape(min(max_bins, n_unique_values) - 1,)
        The increasing numeric values that can be used to separate the bins.
        A given value x will be mapped into bin value i iff
        bining_thresholds[i - 1] < x <= binning_thresholds[i]
    """
    # ignore missing values when computing bin thresholds
    missing_mask = np.isnan(col_data)
    any_missing = missing_mask.any()
    if any_missing:
        col_data = col_data[~missing_mask]

    # If sample_weight is not None and 0-weighted values exist, we need to
    # remove those before calculating the distinct points.
    if sample_weight is not None:
        if any_missing:
            sample_weight = sample_weight[~missing_mask]
        nnz_sw = sample_weight != 0
        col_data = col_data[nnz_sw]
        sample_weight = sample_weight[nnz_sw]

    # The data will be sorted anyway in np.unique and again in percentile, so we do it
    # here. Sorting also returns a contiguous array.
    sort_idx = np.argsort(col_data)
    col_data = col_data[sort_idx]
    if sample_weight is not None:
        sample_weight = sample_weight[sort_idx]

    distinct_values = np.unique(col_data).astype(X_DTYPE)

    if len(distinct_values) == 1:
        return np.asarray([])

    if len(distinct_values) <= max_bins:
        # Calculate midpoints if distinct values <= max_bins
        bin_thresholds = sliding_window_view(distinct_values, 2).mean(axis=1)
    elif sample_weight is None:
        # We compute bin edges using the output of np.percentile with
        # the "averaged_inverted_cdf" interpolation method that is consistent
        # with the code for the sample_weight != None case.
        percentiles = np.linspace(0, 100, num=max_bins + 1)
        percentiles = percentiles[1:-1]
        bin_thresholds = np.percentile(
            col_data, percentiles, method="averaged_inverted_cdf"
        )
        assert bin_thresholds.shape[0] == max_bins - 1
    else:
        percentiles = np.linspace(0, 100, num=max_bins + 1)
        percentiles = percentiles[1:-1]
        bin_thresholds = np.array(