如何在pyspark中创建具有随机值的新列？

Question

我试图用pandas中的随机值初始化新列。我是这样

df['business_vertical'] = np.random.choice(['Retail', 'SME', 'Cor'], df.shape[0])

如何在pyspark中做到这一点？

Answer 1

只需生成一个值列表，然后随机提取它们即可：

from pyspark.sql import functions as F

df.withColumn(
  "business_vertical",
  F.array(
    F.lit("Retail"),
    F.lit("SME"),
    F.lit("Cor"),
  ).getItem(
    (F.rand()*3).cast("int")
  )
)