python - 用列表中的随机项替换 pyspark 数据框列中的某个组

Question

假设数据框如下所示：

ls = [
    ['1', -9.78],
    ['2', 5.38],
    ['1', 8.86],
    ['2', -0.47],
    ['1', -0.19],
    ['1', 4.78],
    ['1', -9.23],
    ['2', -89.32]
]
test = spark.createDataFrame(pd.DataFrame(ls, columns=['col1', 'col2']))
test.show()

输出：

+----+------+
|col1|  col2|
+----+------+
|   1| -9.78|
|   2|  5.38|
|   1|  8.86|
|   2| -0.47|
|   1| -0.19|
|   1|  4.78|
|   1| -9.23|
|   2|-89.32|
+----+------+

我想用从项目列表中随机选择的值替换 col1 == 1 的所有行： ['a', 'b', 'c'] （带替换）。

例如，结果将如下所示：

+----+------+
|col1|  col2|
+----+------+
|   a| -9.78|
|   2|  5.38|
|   a|  8.86|
|   2| -0.47|
|   c| -0.19|
|   b|  4.78|
|   a| -9.23|
|   2|-89.32|
+----+------+

我是 Pyspark 的新手，我很困惑何时使用 when()、where()、withColumn()、select() 等操作。我已经搜索过这个网站，但找不到答案我的问题，所以我真的希望有人能提供帮助！

标签： pythonapache-sparkrandompyspark

Answer 1

您可以定义一个函数以从列表中随机选择一个值。使用pyspark.sql.functions.rand()您可以生成介于 0 和 1 之间的统一随机数。根据所选随机值的值，您可以从列表中选择一个索引。

例如，在您的列表中有 3 个项目的情况下：

• 'a'如果随机数小于 1/3 ，则选择
• 'b'如果随机数小于 2/3 ，则选择
• 否则选择'c'

这是一个可用于将其概括为任何长度选择列表的函数：

from pyspark.sql.functions import rand, when

def pickRandomValue():
    picklist = ['a', 'b', 'c']
    N = len(picklist)
    r = rand()
    c = when(r <= 1.0/N, picklist[0])
    for i in range(1, N-1):
        c = c.when(r <= (i+1.0)/N, picklist[i])
    c = c.otherwise(picklist[-1])
    return c

现在when用来修改col1：

from pyspark.sql.functions import col

test = test.withColumn(
    "col1", 
    when(
        col("col1") == 1,
        pickRandomValue()
    ).otherwise(col("col1").cast("string"))
)
test.show()
#+----+------+
#|col1|  col2|
#+----+------+
#|   c| -9.78|
#|   2|  5.38|
#|   c|  8.86|
#|   2| -0.47|
#|   b| -0.19|
#|   a|  4.78|
#|   a| -9.23|
#|   2|-89.32|
#+----+------+

但是，我必须重申关于 spark 随机数的“谨慎使用”警告。rand()由于 spark 是惰性的，因此每次执行操作时都会重复调用。为了证明这一点，这是第二次调用的结果show()：

test.show()
#+----+------+
#|col1|  col2|
#+----+------+
#|   c| -9.78|
#|   2|  5.38|
#|   b|  8.86|
#|   2| -0.47|
#|   b| -0.19|
#|   b|  4.78|
#|   b| -9.23|
#|   2|-89.32|
#+----+------+

如您所见，结果是不同的。