Factorize on Dask DataFrames

Question

I'm trying to factorize a column in pandas dataframe using the factorize function so that I can have a unique value starting from 0. My question is if there is a way to replicate the same on Dask Dataframes?

Answer 1

Factorization requires a list of unique values, which can be obtained with .unique() , then converting to pandas with .compute() means that we can apply the factorize method:

import pandas as pd
import dask.dataframe as dd

cat = pd.Series(['a', 'a', 'c'])

# calculate uniques
uniques_dask = dd.from_pandas(cat, npartitions=3).unique().compute()

# simple pandas
codes, uniques = pd.factorize(uniques_dask)

# create a mapping
mapping = {k:v for k,v in zip(uniques, codes)}

# apply the mapping
dd.from_pandas(cat, npartitions=3).replace(mapping).compute()