Python 打乱数组数组 pyspark 列

发布于09月05日

我有这样一篇《spark 源》专栏:

                   gm_array
[[1, 4, 6,...], [2, 7, 8,...], [3, 5, 7,...],...]
[[8, 11, 9,...], [7, 2, 6,...], [10, 9, 8,...],...]
[[90, 13, 67,...], [55, 6, 98,...], [1, 6, 2,...],...]
.
.

现在我想把这个数组和数组中的数组混合在一起，然后从前5个数组中选取5个第一个元素.

第一个输出是随机洗牌数组:

                  gm_array
[[19, 6, 1,...], [9, 80, 5,...], [30, 7, 3,...],...]
[[7, 9, 11,...], [6, 8, 7,...], [18, 7, 10,...],...]
[[90, 1, 7,...], [8, 9, 81,...], [6, 5, 1,...],...]
.
.

主数组内前5个数组的第2个输出1个元素:

[19, 9, 30,...]
[7, 6, 18,...]
[[90, 8, 6,...]
.
.

推荐答案

使用一些数组和高阶函数，您可以执行以下操作:

import random
from pyspark.sql import functions as F

# example of input dataframe
df = spark.createDataFrame(
    [
        ([[random.randint(1, 100) for _ in range(5)] for _ in range(6)],)
        for _ in range(4)
    ],
    ["gm_array"]
)

# first step: shuffle arrays
df_shuffled = df.withColumn(
    "gm_array",
    F.shuffle(F.transform("gm_array", lambda x: F.shuffle(x)))
)

# second step: pick top 5 elements
df_top_5 = df_shuffled.withColumn(
    "gm_array",
    F.transform(F.slice("gm_array", 1, 5), lambda x: x[0])
)

df_top_5.show(truncate=False)
#+--------------------+
#|gm_array            |
#+--------------------+
#|[77, 44, 6, 23, 100]|
#|[40, 57, 10, 32, 27]|
#|[3, 45, 17, 9, 9]   |
#|[62, 39, 10, 95, 17]|
#+--------------------+