Spark：遍历每行中的列以创建新的数据表

Question

假设我有这样一个数据frame：

+-----------+-----------+-----------+-----------+------------+--+
|   ColA    |   ColB    |   ColC    |   ColD    |    ColE    |  |
+-----------+-----------+-----------+-----------+------------+--+
| ''        | sample_1x | sample_1y | ''        | sample_1z  |  |
| sample2_x | sample2_y | ''        | ''        | ''         |  |
| sample3_x | ''        | ''        | ''        | sample3_y  |  |
| sample4_x | sample4_y | ''        | sample4_z | sample4_zz |  |
| sample5_x | ''        | ''        | ''        | ''         |  |
+-----------+-----------+-----------+-----------+------------+--+

我希望创建另一个dataframe，在每行中从左到右显示关系，同时跳过具有空值的列。此外，只有1个有效列记录的行将被排除。例如：

+-----------+------------+-----------+
|   From    |     To     |   Label   |
+-----------+------------+-----------+
| sample1_x | sample1_y  | ColB_ColC |
| sample1_y | sample1_z  | ColC_ColE |
| sample2_x | sample2_y  | ColA_ColB |
| sample3_x | sample3_y  | ColA_ColE |
| sample4_x | sample4_y  | ColA_ColB |
| sample4_y | sample4_z  | ColB_ColD |
| sample4_z | sample4_zz | ColD_ColE |
+-----------+------------+-----------+

我认为方法应该是编写一个包含此逻辑的UDF，但我不完全确定如何返回一个全新的DF，因为我习惯了UDF只是在同一个DF中创建另一列。或者是否有另一个spark函数可以比创建UDF更容易地处理这种情况？如果有关系，就使用pyspark。

Answer 1

您可以使用接受数组参数并返回结构数组的udf，例如：

from pyspark.sql import functions as F

df.show()
+---------+---------+---------+---------+----------+
|     ColA|     ColB|     ColC|     ColD|      ColE|
+---------+---------+---------+---------+----------+
|     null|sample_1x|sample_1y|     null| sample_1z|
|sample2_x|sample2_y|     null|     null|      null|
|sample3_x|     null|     null|     null| sample3_y|
|sample4_x|sample4_y|     null|sample4_z|sample4_zz|
|sample5_x|     null|     null|     null|      null|
+---------+---------+---------+---------+----------+

# columns that get involved, will group them into an array using F.array(cols)
cols = df.columns

# defind function to convert array into array of structs
def find_route(arr, cols):
    d = [ (cols[i],e) for i,e in enumerate(arr) if e is not None ]
    return [ {'From':d[i][1], 'To':d[i+1][1], 'Label':d[i][0]+'_'+d[i+1][0]} for i in range(len(d)-1) ]

# set up the UDF and add cols as an extra argument
udf_find_route = F.udf(lambda a: find_route(a, cols), 'array<struct<From:string,To:string,Label:string>>')

# retrive the data from the array of structs after array-explode
df.select(F.explode(udf_find_route(F.array(cols))).alias('c1')).select('c1.*').show()
+---------+----------+---------+
|     From|        To|    Label|
+---------+----------+---------+
|sample_1x| sample_1y|ColB_ColC|
|sample_1y| sample_1z|ColC_ColE|
|sample2_x| sample2_y|ColA_ColB|
|sample3_x| sample3_y|ColA_ColE|
|sample4_x| sample4_y|ColA_ColB|
|sample4_y| sample4_z|ColB_ColD|
|sample4_z|sample4_zz|ColD_ColE|
+---------+----------+---------+