Skip to content

Commit

Permalink
Reduce memory usage of as_categorical_column (rapidsai#14138)
Browse files Browse the repository at this point in the history
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

```python
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

![Screenshot from 2023-09-20 14-49-27](https://github.com/rapidsai/cudf/assets/1126981/08782501-c233-4efd-b4d6-a378cea82a82)

### After

![Screenshot from 2023-09-20 14-49-42](https://github.com/rapidsai/cudf/assets/1126981/93193bfb-a93e-45bf-8e5a-24289efc77c4)

Authors:
  - Lawrence Mitchell (https://github.com/wence-)

Approvers:
  - GALI PREM SAGAR (https://github.com/galipremsagar)
  - Bradley Dice (https://github.com/bdice)
  - Ashwin Srinath (https://github.com/shwina)

URL: rapidsai#14138
  • Loading branch information
wence- authored Sep 20, 2023
1 parent 40d4cc5 commit e87d2fc
Showing 1 changed file with 8 additions and 9 deletions.
17 changes: 8 additions & 9 deletions python/cudf/cudf/core/column/column.py
Original file line number Diff line number Diff line change
Expand Up @@ -1390,20 +1390,19 @@ def _return_sentinel_column():
except ValueError:
return _return_sentinel_column()

codes = arange(len(cats), dtype=dtype)
left_gather_map, right_gather_map = cpp_join(
[self], [cats], how="left"
)
codes = codes.take(
right_gather_map, nullify=True, check_bounds=False
).fillna(na_sentinel.value)

codes = libcudf.copying.gather(
[arange(len(cats), dtype=dtype)], right_gather_map, nullify=True
)
del right_gather_map
# reorder `codes` so that its values correspond to the
# values of `self`:
order = arange(len(self))
order = order.take(left_gather_map, check_bounds=False).argsort()
codes = codes.take(order)
return codes
(codes,) = libcudf.sort.sort_by_key(
codes, [left_gather_map], [True], ["last"], stable=True
)
return codes.fillna(na_sentinel.value)


def column_empty_like(
Expand Down

0 comments on commit e87d2fc

Please sign in to comment.