在数组字段（可变长度）上应用 Udf 并将其拆分为 pyspark 中的列

Question

我想对长度为 可变 (0-4000) 的数组字段应用逻辑并将其拆分为列。具有爆炸、创建新列和重命名列的 udf 可以完成这项工作，但我不确定如何将其作为 udf 迭代应用。 UDF 将采用可变长度数组字段并将一组新列 (0-4000) 返回到数据帧。如下所示的示例输入数据框

+--------------------+--------------------+
|             hashval|    dec_spec (array|
+--------------------+--------------------+
|3c65252a67546832d...|[8.02337424829602...|
|f5448c29403c80ea7...|[7.50372884795069...|
|94ff32cd2cfab9919...|[5.85195317398756...|
+--------------------+--------------------+

输出应该是这样的

+--------------------+--------------------+
    |             hashval|    dec_spec (array|   ftr_1    |    ftr_2 | ftr_3 |...
    +--------------------+--------------------+-----------+---------+--------+
    |3c65252a67546832d...|[8.02337424829602...|  8.023   | 3.21       | 4.23.....
    |f5448c29403c80ea7...|[7.50372884795069...| 7.502    | 8.23       |2.125
    |94ff32cd2cfab9919...|[5.85195317398756...|
    +--------------------+--------------------+

udf 可以采用如下的一些逻辑

df_grp = df2.withColumn("explode_col", F.explode_outer("dec_spec"))
df_grp = df_grp.groupBy("hashval").pivot("explode_col").agg(F.avg("explode_col"))

下面用于重命名列

count = 1
for col in df_grp.columns:
  if col != "hashval":
    df_grp = df_grp.withColumnRenamed(col, "ftr"+str(count))
    count = count+1

感谢任何帮助。

PS 上面的代码，在这里得到了论坛其他人的帮助。

Answer 1

模拟样本数据

from pyspark.sql import functions as sf
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType, StructType, StructField, StringType
sdf1 = sc.parallelize([["aaa", "1,2,3"],["bbb", "1,2,3,4,5"]]).toDF(["hash_val", "arr_str"])
sdf2 = sdf1.withColumn("arr", sf.split("arr_str", ","))
sdf2.show()

+--------+---------+---------------+
|hash_val|  arr_str|            arr|
+--------+---------+---------------+
|     aaa|    1,2,3|      [1, 2, 3]|
|     bbb|1,2,3,4,5|[1, 2, 3, 4, 5]|
+--------+---------+---------------+

udf 使所有数组长度相同

schema = ArrayType(StringType())

def fill_list(input_list, input_length):
    fill_len = input_length - len(input_list)
    if fill_len > 0:
        input_list += [None]*(fill_len)
    
    return input_list[0:input_length]

fill_list_udf = udf(fill_list, schema)
sdf3 = sdf2.withColumn("arr1", fill_list_udf(sf.col("arr"), sf.lit(3)))
sdf3.show()

+--------+---------+---------------+---------+
|hash_val|  arr_str|            arr|     arr1|
+--------+---------+---------------+---------+
|     aaa|    1,2,3|      [1, 2, 3]|[1, 2, 3]|
|     bbb|1,2,3,4,5|[1, 2, 3, 4, 5]|[1, 2, 3]|
+--------+---------+---------------+---------+

展开它们

sdf3.select("hash_val", *[sf.col("arr1")[i] for i in range(3)]).show()
+--------+-------+-------+-------+
|hash_val|arr1[0]|arr1[1]|arr1[2]|
+--------+-------+-------+-------+
|     aaa|      1|      2|      3|
|     bbb|      1|      2|      3|
+--------+-------+-------+-------+