Reduce memory usage of as_categorical_column

[wence-]

Description

The main culprit is in the way the codes returned from _label_encoding were being ordered. We were generating an int64 column for the order, gathering through the left gather map, and then argsorting, before using that ordering as a gather map for the codes.

We note that gather(y, with=argsort(x)) is equivalent to sort_by_key(y, with=x) so use that instead (avoiding an unnecessary gather). Furthermore we also note that gather([0..n), with=x) is just equivalent to x, so we can avoid a gather too.

This reduces the peak memory footprint of categorifying a random column of 500_000_000 int32 values where there are 100 unique values from 24.75 GiB to 11.67 GiB.

Test code

import cudf
import cupy as cp

K = 100
N = 500_000_000
rng = cp.random._generator.RandomState()
column = cudf.core.column.as_column(rng.choice(cp.arange(K, dtype="int32"), size=(N,), replace=True))
column = column.astype("category", ordered=False)

Before

After

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[bdice]

Great! Note that this is an example of the performance antipattern discussed in #13557.

[wence-]

/merge

[harrism]

Is performance affected?

[wence-]

Is performance affected?

Yes, but positively, I run:

import time
import cupy as cp
import cudf
import rmm

rmm.reinitialize(pool_allocator=True)

rng = cp.random._generator.RandomState(seed=108)
for K in [2**4, 2**10, 2**12, 2**14, 2**16]:
    for N in [1_000_000, 10_000_000, 100_000_000, 250_000_000]:
        col = cudf.core.column.as_column(rng.choice(cp.arange(K, dtype="uint32"), size=N, replace=True))
        start = time.time()
        for _ in range((reps := 1_000_000_000 // N)):
            y = col.astype("category", ordered=False)
            del y
        end = time.time()
        del col

Across column sizes and number of unique values, the new code is between 25 and 30% faster.

[harrism]

Excellent!