Не уверен, что это то, что вы ищете, но вот как я это сделал:
import random
from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType
npart = 50000
def rand(n = npart):
return random.randint(1, n)
salt = udf(rand, IntegerType())
DF.withColumn("salt", salt())