在 pyspark 中添加带有压缩列表的新列作为常量值

Posted 2023-04-15

【中文标题】在 pyspark 中添加带有压缩列表的新列作为常量值

【英文标题】：Add new column with zipped list as a constant value in pyspark

【发布时间】：2020-04-06 03:55:20

【问题描述】：

我有 payspark 数据框，我想添加具有恒定值 x 的新列，它是压缩列表：

x = [('1', 'hello'),('2', 'Hi'),('3', 'Hello')]

但是当我运行这段代码时：

df = df.withColumn('case', x)

我收到此错误：

AssertionError: col should be Column

我怎样才能给这个列表提供结构来处理这个错误，我知道对于 int 或字符串值我们可以使用 Lit 函数，但是对于这种格式我不知道该怎么做。

【参考方案1】：

您可以将 string literals 放入 structs 到 array。

x = [('1', 'hello'),('2','Hi'),('3', 'Hello')]

df.withColumn("col1", F.array(*[F.struct(F.lit(i[0]),F.lit(i[1])) for i in x])).show(truncate=False)

+---------------------------------+
|col1                             |
+---------------------------------+
|[[1, hello], [2, Hi], [3, Hello]]|
|[[1, hello], [2, Hi], [3, Hello]]|
+---------------------------------+

Or，您创建 array 的 arrays：

x = [('1', 'hello'),('2','Hi'),('3', 'Hello')]

df.withColumn("col1", F.array(*[F.array(F.lit(i[0]),F.lit(i[1])) for i in x])).show(truncate=False)

+---------------------------------+
|col1                             |
+---------------------------------+
|[[1, hello], [2, Hi], [3, Hello]]|
|[[1, hello], [2, Hi], [3, Hello]]|
+---------------------------------+

Or，您可以以类似的方式使用 create_map：

x = [('1', 'hello'),('2','Hi'),('3', 'Hello')]

df.withColumn("col1",F.array(*[F.create_map(F.lit(i[0]),F.lit(i[1])) for i in x])).show(truncate=False)

+---------------------------------------+
|col1                                   |
+---------------------------------------+
|[[1 -> hello], [2 -> Hi], [3 -> Hello]]|
|[[1 -> hello], [2 -> Hi], [3 -> Hello]]|
+---------------------------------------+

【讨论】：

效果很好，但是对于 pandas 数据框，我使用此列来应用功能，但在 pyspark 中它有问题，也许我应该更改我的 udf 中的输入类型，谢谢您的帮助