Repository: dask/dask
File: dask/dataframe/partitionquantiles.py (Python)

Code before:
    """
    from dask.array.dispatch import percentile_lookup as _percentile
    from dask.array.utils import array_safe

    length = len(df)
    if length == 0:
        return ()
    random_state = np.random.RandomState(state)
    qs = sample_percentiles(num_old, num_new, length, upsample, random_state)
    data = df
    interpolation = "linear"

    if isinstance(data.dtype, pd.CategoricalDtype):
        data = data.cat.codes
        interpolation = "nearest"
    elif is_datetime64_dtype(data.dtype) or is_integer_dtype(data.dtype):
        interpolation = "nearest"

    # FIXME: pandas quantile doesn't work with some data types (e.g. strings).
    # We fall back to an ndarray as a workaround.
    try:
        vals = data.quantile(q=qs / 100, interpolation=interpolation).values
    except (TypeError, NotImplementedError):
        try:
            vals, _ = _percentile(array_safe(data, like=data.values), qs, interpolation)
        except (TypeError, NotImplementedError):
            # `data.values` doesn't work for cudf, so we need to
            # use `quantile(..., method="table")` as a fallback
            interpolation = "nearest"
            vals = (
                data.to_frame()
                .quantile(
                    q=qs / 100,
                    interpolation=interpolation,
                    numeric_only=False,
                    method="table",
                )
                .iloc[:, 0]
            )

    if (
        is_cupy_type(data)
        and interpolation == "linear"
        and np.issubdtype(data.dtype, np.integer)
    ):
        vals = np.round(vals).astype(data.dtype)
        if qs[0] == 0:
            # Ensure the 0th quantile is the minimum value of the data
            vals[0] = data.min()
    vals_and_weights = percentiles_to_weights(qs, vals, length)
    return vals_and_weights

Review comment (suggestion):
Hopefully this will ensure that the unnecessary `_percentile` code path gets removed when pandas<1.5 is no longer supported.

Code after:
        Scale factor to increase the number of percentiles calculated in
        each partition.  Use to improve accuracy.
    """
    from dask.array.dispatch import percentile_lookup as _percentile
    from dask.array.utils import array_safe

    length = len(df)
    if length == 0:
        return ()
    random_state = np.random.RandomState(state)
    qs = sample_percentiles(num_old, num_new, length, upsample, random_state)
    data = df
    interpolation = "linear"

    if isinstance(data.dtype, pd.CategoricalDtype):
        data = data.cat.codes
        interpolation = "nearest"
    elif is_datetime64_dtype(data.dtype) or is_integer_dtype(data.dtype):
        interpolation = "nearest"

    try:
        # Plan A: Use `Series.quantile`
        vals = data.quantile(q=qs / 100, interpolation=interpolation)
    except (TypeError, NotImplementedError):
        # Series.quantile doesn't work with some data types (e.g. strings)
        if PANDAS_GE_150:
            # Fall back to DataFrame.quantile with "nearest" interpolation
            interpolation = "nearest"
            vals = (
                data.to_frame()
                .quantile(
                    q=qs / 100,
                    interpolation=interpolation,
                    numeric_only=False,
                    method="table",
                )
                .iloc[:, 0]
            )
        else:
            # Fall back to ndarray (Not supported for cuDF)
            vals, _ = _percentile(array_safe(data, like=data.values), qs, interpolation)

    # Convert to array if necessary (and possible)
    if is_series_like(vals):
        try:
            vals = vals.values
        except (ValueError, TypeError):
            # cudf->cupy won't work if nulls are present,
            # or if this is a string dtype
            pass