来自ArrayType Pyspark列的随机样本

Question

我在Pyspark数据框中有一个具有类似结构的列

Column1
[a,b,c,d,e]
[c,b,d,f,g,h,i,p,l,m]

我想返回另一列，其中每行中的每个数组都是随机选择的，函数中指定的数量。

所以类似data.withColumn("sample", SOME_FUNCTION("column1", 5))返回：

sample
[a,b,c,d,e]
[c,b,h,i,p]

希望避免使用Python UDF，感觉应该有可用的功能？

此作品：

import random
def random_sample(population):
    return(random.sample(population, 5))

udf_random = F.udf(random_sample, T.ArrayType(T.StringType()))
df.withColumn("sample", udf_random("column1")).show()

但是正如我所说，最好避免使用UDF。

Answer 1

对于spark 2.4+，请使用shuffle和shuffle：

slice