Python 如何获取Pyspark Dataframe中另一列中给定了多列值的列表列？

有没有办法创建一个新的列，如Pyspark中下面显示的Dataframe

我一直在尝试列表理解：

import pyspark.functions as F

df.withColumn('result', [F.col(colname) for colname in F.col('colList')])

但是不起作用

预期结果是：

+----+----+----+----+----+---------------+------+
|col1|col2|col3|col4|col5|        colList|result|
+----+----+----+----+----+---------------+------+
|   1|   2|   0|   3|   4|['col1','col2']| [1,2]|
|   1|   2|   0|   3|   4|['col2','col3']| [2,0]|
|   1|   2|   0|   3|   4|['col1','col3']| [1,0]|
|   1|   2|   0|   3|   4|['col3','col4']| [0,3]|
|   1|   2|   0|   3|   4|['col2','col5']| [2,4]|
|   1|   2|   0|   3|   4|['col4','col5']| [3,4]|
+----+----+----+----+----+---------------+------+

作为下一步，我们为数组colList中的各个列创建列

df = df.withColumn('first_col',col('colList')[0])
df = df.withColumn('second_col',col('colList')[1])
df.show()
+----+----+----+----+----+------------+---------+----------+
|col1|col2|col3|col4|col5|     colList|first_col|second_col|
+----+----+----+----+----+------------+---------+----------+
|   1|   2|   0|   3|   4|[col1, col2]|     col1|      col2|
|   1|   2|   0|   3|   4|[col2, col3]|     col2|      col3|
|   1|   2|   0|   3|   4|[col1, col3]|     col1|      col3|
|   1|   2|   0|   3|   4|[col3, col4]|     col3|      col4|
|   1|   2|   0|   3|   4|[col2, col5]|     col2|      col5|
|   1|   2|   0|   3|   4|[col4, col5]|     col4|      col5|
+----+----+----+----+----+------------+---------+----------+

具有整数值的列的列表-

concerned_columns = [x for x in df.columns if x not in {'colList','first_col','second_col'}]
print(concerned_columns)
    ['col1', 'col2', 'col3', 'col4', 'col5']

现在，最重要的部分是，我们使用spark 2.+之后的函数创建列名与其各自值之间的映射

# Maping - (column name, column values)
col_name_value_mapping = create_map(*chain.from_iterable(
    (lit(c), col(c)) for c in concerned_columns
))

最后，应用此映射获取存储在第一列和第二列中的列的值，并使用将它们放入数组中

---

如果您使用的是Spark 2.40+，可以尝试使用

transform（）

函数来简化任务。

# Maping - (column name, column values)
col_name_value_mapping = create_map(*chain.from_iterable(
    (lit(c), col(c)) for c in concerned_columns
))

df = df.withColumn('result', struct(col_name_value_mapping[col('first_col')],col_name_value_mapping[col('second_col')]))
df = df.drop('first_col','second_col')
df.show()
+----+----+----+----+----+------------+------+
|col1|col2|col3|col4|col5|     colList|result|
+----+----+----+----+----+------------+------+
|   1|   2|   0|   3|   4|[col1, col2]| [1,2]|
|   1|   2|   0|   3|   4|[col2, col3]| [2,0]|
|   1|   2|   0|   3|   4|[col1, col3]| [1,0]|
|   1|   2|   0|   3|   4|[col3, col4]| [0,3]|
|   1|   2|   0|   3|   4|[col2, col5]| [2,4]|
|   1|   2|   0|   3|   4|[col4, col5]| [3,4]|
+----+----+----+----+----+------------+------+