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

Code before:
        q_weights = np.cumsum(trimmed_weights)
        q_target = np.linspace(0, q_weights[-1], trimmed_npartitions + 1)

        left = np.searchsorted(q_weights, q_target, side="left")
        right = np.searchsorted(q_weights, q_target, side="right") - 1
        # stay inbounds
        np.maximum(right, 0, right)
        lower = np.minimum(left, right)
        trimmed = trimmed_vals[lower]

        rv = np.concatenate([trimmed, jumbo_vals])
        rv.sort()

    if is_categorical_dtype(dtype):
        rv = pd.Categorical.from_codes(rv, info[0], info[1])
    elif is_datetime64tz_dtype(dtype):
        rv = pd.DatetimeIndex(rv).tz_localize(dtype.tz)
    elif "datetime64" in str(dtype):
        rv = pd.DatetimeIndex(rv, dtype=dtype)
    elif rv.dtype != dtype:
        rv = rv.astype(dtype)
    return rv


@methods.percentiles_summary.register(pd.Series)
def percentiles_summary_pandas(df, num_old, num_new, upsample, state):
    """Summarize data using percentiles and derived weights.

    These summaries can be merged, compressed, and converted back into
    approximate percentiles.

    Parameters
    ----------
    df: pandas.Series
        Data to summarize
    num_old: int
        Number of partitions of the current object
    num_new: int
        Number of partitions of the new object
    upsample: float
        Scale factor to increase the number of percentiles calculated in
        each partition.  Use to improve accuracy.
    """
    from dask.array.percentile import _percentile

    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.values

Review comment (performance):
Is it possible to avoid dispatching this one? 

If we added a `tolist()` method to cudf, and allowed for a very small amount of host-device transfer for some random values, would we need to split this off?  Dispatch functions open up some code complexity.  I would prefer to use them sparingly if we can avoid them.  Obviously in some cases we can't avoid them, but this seems like a case where I'm not sure.  Thoughts?

Code after:
        right = np.searchsorted(q_weights, q_target, side="right") - 1
        # stay inbounds
        np.maximum(right, 0, right)
        lower = np.minimum(left, right)
        trimmed = trimmed_vals[lower]

        rv = np.concatenate([trimmed, jumbo_vals])
        rv.sort()

    if is_categorical_dtype(dtype):
        rv = pd.Categorical.from_codes(rv, info[0], info[1])
    elif is_datetime64tz_dtype(dtype):
        rv = pd.DatetimeIndex(rv).tz_localize(dtype.tz)
    elif "datetime64" in str(dtype):
        rv = pd.DatetimeIndex(rv, dtype=dtype)
    elif rv.dtype != dtype:
        rv = rv.astype(dtype)
    return rv


def percentiles_summary(df, num_old, num_new, upsample, state):
    """Summarize data using percentiles and derived weights.

    These summaries can be merged, compressed, and converted back into
    approximate percentiles.

    Parameters
    ----------
    df: pandas.Series
        Data to summarize
    num_old: int
        Number of partitions of the current object
    num_new: int
        Number of partitions of the new object
    upsample: float
        Scale factor to increase the number of percentiles calculated in
        each partition.  Use to improve accuracy.
    """
    from dask.array.percentile import _percentile

    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.values
    interpolation = "linear"
    if is_categorical_dtype(data):
        data = data.codes
        interpolation = "nearest"
    vals, n = _percentile(data, qs, interpolation=interpolation)