如何在pyspark中压缩/连接值和列表

Question

我正在处理一个使用 4 个输入的函数。

为此，我想得到一个总结这 4 个元素的列表。但是我有两个变量，其中数据是唯一的，两个变量由列表组成。我可以用 压缩这两个列表arrays_zip，但无法获得包含 4 个元素的数组列表：

+----+----+---------+---------+
| l1 | l2 |   l3    |   l4    |
+----+----+---------+---------+
| 1  | 5  | [1,2,3] | [2,2,2] |
| 2  | 9  | [8,2,7] | [1,7,7] |
| 3  | 3  | [8,4,9] | [5,1,3] |
| 4  | 1  | [5,5,3] | [8,4,3] |

我想得到什么：

+----+----+---------+---------+------------------------------------------+
| l1 | l2 |   l3    |   l4    |                       l5                 |
+----+----+---------+---------+------------------------------------------+
| 1  | 5  | [1,2,3] | [2,2,2] | [[1, 5, 1, 2],[1, 5, 2, 2],[1, 5, 3, 2]] |
| 2  | 9  | [8,2,7] | [1,7,7] | [[2, 9, 8, 1],[2, 9, 2, 7],[2, 9, 7, 7]] |
| 3  | 3  | [8,4,9] | [5,1,3] | [[3, 3, 8, 5],[3 ,3, 4, 1],[3, 3, 9, 3]] |
| 4  | 1  | [5,5,3] | [8,4,3] | [[4, 1, 5, 8],[4, 1, 5, 4],[4, 1, 3, 3]] |

我的想法是用 l3 大小转换列表中的 l1 和 l2，然后应用arrays_zip. 我没有找到创建此列表的一致方法。

只要我获得了这个列表列表，我就会应用如下函数：

def is_good(data):
  a,b,c,d = data
  return a+b+c+d

is_good_udf = f.udf(lambda x: is_good(x), ArrayType(FloatType()))

spark.udf.register("is_good_udf ", is_good, T.FloatType())

我的猜测是构建这样的东西，感谢@kafels，其中对于列表的每一行和每个列表，它应用函数：

df.withColumn("tot", f.expr("transform(l5, y -> is_good_udf(y))"))

例如，为了获得[9, 10, 11]第一行的结果列表。

Answer 1

您可以使用expr函数并应用TRANSFORM：

import pyspark.sql.functions as f


df = df.withColumn('l5', f.expr("""TRANSFORM(arrays_zip(l3, l4), el -> array(l1, l2, el.l3, el.l4))"""))

# +---+---+---------+---------+------------------------------------------+
# |l1 |l2 |l3       |l4       |l5                                        |
# +---+---+---------+---------+------------------------------------------+
# |1  |5  |[1, 2, 3]|[2, 2, 2]|[[1, 5, 1, 2], [1, 5, 2, 2], [1, 5, 3, 2]]|
# |2  |9  |[8, 2, 7]|[1, 7, 7]|[[2, 9, 8, 1], [2, 9, 2, 7], [2, 9, 7, 7]]|
# |3  |3  |[8, 4, 9]|[5, 1, 3]|[[3, 3, 8, 5], [3, 3, 4, 1], [3, 3, 9, 3]]|
# |4  |1  |[5, 5, 3]|[8, 4, 3]|[[4, 1, 5, 8], [4, 1, 5, 4], [4, 1, 3, 3]]|
# +---+---+---------+---------+------------------------------------------+