apache-spark - Pyspark concat 列基于其他列值

Question

我正在使用 pyspark 下面是我的数据框

col_list   | col_1| col_2 |...|col_n
col_1,col2 | x    | y     |...|n
col_2,col_n| a    | b     |...|z

我想要我的 col_list 列值中存在的 concat 列

预期的DF：

col_list   | col_1| col_2 |...|col_n | concat_cols
col_1,col2 | x    | y     |...|n     | [col_1:x,col_2:y]
col_2,col_n| a    | b     |...|z     | [col_2:b,col_n:z]

尝试以下方法：

1. 创建 udf 以传递 col_list 并创建表达式以创建数组 col 像 array(df['col_1'],df['col_2']) ，但这将创建带有 string 的列，不确定如何执行表达式

2. 使用 create map 在循环中传递 col_list 并创建列，需要对此进行检查。

标签： apache-sparkpysparkapache-spark-sql

Answer 1

您可以使用一些地图方法：

import pyspark.sql.functions as F

df2 = df.withColumn(
    'col_map', 
    F.map_from_arrays(
        F.array([F.lit(c) for c in df.columns[1:]]), 
        F.array(df.columns[1:])
    )
).withColumn(
    'concat_cols', 
    F.expr("""
        map_from_arrays(
            split(col_list, ','), 
            transform(split(col_list, ','), x -> col_map[x])
        )
    """)
).drop('col_map')

df2.show(truncate=False)
+-----------+-----+-----+-----+------------------------+
|col_list   |col_1|col_2|col_n|concat_cols             |
+-----------+-----+-----+-----+------------------------+
|col_1,col_2|x    |y    |n    |[col_1 -> x, col_2 -> y]|
|col_2,col_n|a    |b    |z    |[col_2 -> b, col_n -> z]|
+-----------+-----+-----+-----+------------------------+

如果您更喜欢字符串表示而不是地图类型列，您可以这样做

import pyspark.sql.functions as F

df2 = df.withColumn(
    'col_map', 
    F.map_from_arrays(
        F.array([F.lit(c) for c in df.columns[1:]]), 
        F.array(df.columns[1:])
    )
).withColumn(
    'concat_cols', 
    F.expr("""
        concat_ws(',', 
            transform(split(col_list, ','), x -> concat_ws(':', x, col_map[x]))
        )
    """)
).drop('col_map')

df2.show(truncate=False)
+-----------+-----+-----+-----+---------------+
|col_list   |col_1|col_2|col_n|concat_cols    |
+-----------+-----+-----+-----+---------------+
|col_1,col_2|x    |y    |n    |col_1:x,col_2:y|
|col_2,col_n|a    |b    |z    |col_2:b,col_n:z|
+-----------+-----+-----+-----+---------------+