在Pyspark中使用随机固定值向现有数据框添加列

Question

我是Pyspark的新手，我正在尝试为现有数据框添加新列。新列应仅包含4个固定值（例如1,2,3,4），并且我想随机选择每行的一个值。

我该怎么做？

Answer 1

Pyspark数据帧是不可变的，因此您必须返回一个新数据帧（例如，您无法以Pandas数据帧的方式分配它）。要做你想做的事，请使用udf：

from pyspark.sql.functions import udf
import numpy as np

df = <original df>

udf_randint = udf(np.random.randint(1, 4))
df_new = df.withColumn("random_num": udf_randint)