python - 如何在 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")
  )
)