使用pyspark将两个字符串列值连接到整数数组

Posted 2023-04-15

【中文标题】使用pyspark将两个字符串列值连接到整数数组

【英文标题】：Concatenate two string column values to array of integer using pyspark

【发布时间】：2021-06-17 10:43:54

【问题描述】：

我有一个具有不同事件类型（10、11、12）的数据框，数据因事件类型而异，我的要求是，当事件类型 = 11 时，我需要将两列数据（传感器，字段 1）连接到数组[0,1] 在传感器列中并将数据向左移动，我已经给出了我想要的输出，有没有一种方法可以使用 pyspark 实现这一点。

我能够使用 spark.sql concat() 函数进行连接，但正在寻找更通用的代码，该代码将连接列值并将其向左移动并将最后一列值分配为 NULL。

eventtype   sensor  field1  field2  eventcategory
 11           1       0     SE_500  500

预期输出：

eventtype, sensor, field1, field2, eventcategory
 11         [1,0]   SE_500   500    NULL

【参考方案1】：

您可以使用when 来构建适当的选择。请注意，一列只能有一种类型，因此如果 eventtype != 11，则必须将传感器列转换为一个元素的数组。

import pyspark.sql.functions as F

df2 = df.select(
    'eventtype', 
    F.when(
        F.col('eventtype') == 11, 
        F.array('sensor', 'field1')
    ).otherwise(F.array('sensor')).alias('sensor'), 
    *[
        F.when(
            F.col('eventtype') == 11, 
            F.col(df.columns[i+3]) if i+3 < len(df.columns) else F.lit(None)
        ).otherwise(F.col(c)).alias(c) 
        for (i, c) in enumerate(df.columns[2:])
    ]
)

df2.show()
+---------+------+------+------+-------------+
|eventtype|sensor|field1|field2|eventcategory|
+---------+------+------+------+-------------+
|       11|[1, 0]|SE_500|   500|         null|
|       10|   [1]|     0|SE_500|          500|
+---------+------+------+------+-------------+