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

Question

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

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

我如何在 pyspark 中执行此操作？

原文由 subash poudel 发布，翻译遵循 CC BY-SA 4.0 许可协议

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")
  )
)

原文由 Steven 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

以下是使用 quinn 中的 array_choice 函数解决此问题的方法：

 import quinn

df = spark.createDataFrame([('a',), ('b',), ('c',)], ['letter'])
cols = list(map(lambda c: F.lit(c), ['Retail', 'SME', 'Cor']))
df.withColumn('business_vertical', quinn.array_choice(F.array(cols))).show()

 +------+-----------------+
|letter|business_vertical|
+------+-----------------+
|     a|              SME|
|     b|           Retail|
|     c|              SME|
+------+-----------------+

array_choice 是通用的，可轻松用于从现有 ArrayType 列中选择随机值。假设您有以下 DataFrame。

 +------------+
|     letters|
+------------+
|   [a, b, c]|
|[a, b, c, d]|
|         [x]|
|          []|
+------------+

这是您可以随机获取一封信的方法。

 actual_df = df.withColumn(
    "random_letter",
    quinn.array_choice(F.col("letters"))
)
actual_df.show()

 +------------+-------------+
|     letters|random_letter|
+------------+-------------+
|   [a, b, c]|            a|
|[a, b, c, d]|            d|
|         [x]|            x|
|          []|         null|
+------------+-------------+

这是 array_choice 函数定义：

 def array_choice(col):
    index = (F.rand()*F.size(col)).cast("int")
    return col[index]

这篇 文章 更详细地解释了从 PySpark 数组中获取随机值。

原文由 Powers 发布，翻译遵循 CC BY-SA 4.0 许可协议