将带有 int 标志的列转换为 pyspark 中的字符串数组

Question

我有一个数据框，其中有一列称为“特征”，它是一个由多个标志组成的整数。

我需要将此列转换为字符串列表（用于弹性搜索索引）。转换如下所示。

TRAIT_0 = 0
TRAIT_1 = 1
TRAIT_2 = 2
def flag_to_list(flag: int) -> List[str]:
    trait_list = []
    if flag & (1 << TRAIT_0):
        trait_list.append("TRAIT_0")
    elif flag & (1 << TRAIT_1):
        trait_list.append("TRAIT_1")
    elif flag & (1 << TRAIT_2):
        trait_list.append("TRAIT_2")

    return trait_list

在 pyspark 中进行这种转换的最有效方法是什么？我看到很多关于如何进行字符串连接和拆分的示例，但不是这样的操作。

使用 pyspark 版本 2.4.5

输入 json 如下所示： { "name": "John Doe", "traits": 5 } 输出 json 应如下所示： { "name": "John Doe", "traits": ["TRAIT_0", "TRAIT_2"] }

Answer 1

IIUC，你可以试试 SparkSQL 的内置函数： (1) 使用conv + split 转换整数(base-10) -> 二进制(base-2) -> 字符串 -> 字符串数组(颠倒），（2）基于0或1的值及其数组索引到filter和transform将数组放入对应的命名特征数组中：

from pyspark.sql.functions import expr

df = spark.createDataFrame([("name1", 5),("name2", 1),("name3", 0),("name4", 12)], ['name', 'traits'])
#DataFrame[name: string, traits: bigint]

traits = [ "Traits_{}".format(i) for i in range(8) ]
traits_array = "array({})".format(",".join("'{}'".format(e) for e in traits))
# array('Traits_0','Traits_1','Traits_2','Traits_3','Traits_4','Traits_5','Traits_6','Traits_7')

sql_expr = """

    filter(
      transform(
        /* convert int -> binary -> string -> array of strings, and then reverse the array */
        reverse(split(string(conv(traits,10,2)),'(?!$)')),
        /* take the corresponding items from the traits_array when value > 0, else NULL */
        (x,i) -> {}[IF(x='1',i,NULL)]
      ),
      /* filter out NULL items from the array */
      y -> y is not NULL
    ) AS trait_list

""".format(traits_array)
# filter(
#   transform(
#     reverse(split(string(conv(traits,10,2)),'(?!$)')),
#     (x,i) -> array('Traits_0','Traits_1','Traits_2','Traits_3','Traits_4','Traits_5','Traits_6','Traits_7')[IF(x='1',i,NULL)]
#   ),
#   y -> y is not NULL
# )

df.withColumn("traits_list", expr(sql_expr)).show(truncate=False)
+-----+------+--------------------+
|name |traits|traits_list         |
+-----+------+--------------------+
|name1|5     |[Traits_0, Traits_2]|
|name2|1     |[Traits_0]          |
|name3|0     |[]                  |
|name4|12    |[Traits_2, Traits_3]|
+-----+------+--------------------+

以下是运行reverse(split(string(conv(traits,10,2)),'(?!$)')) 后的结果，注意拆分模式(?!$) 用于避免显示为最后一个数组项的NULL。

df.selectExpr("*", "reverse(split(string(conv(traits,10,2)),'(?!$)')) as t1").show()
+-----+------+------------+
| name|traits|          t1|
+-----+------+------------+
|name1|     5|   [1, 0, 1]|
|name2|     1|         [1]|
|name3|     0|         [0]|
|name4|    12|[0, 0, 1, 1]|
+-----+------+------------+

Answer 2

我们可以定义一个 UDF 来包装你的函数，然后调用它。这是一些示例代码：

from typing import List
from pyspark.sql.types import ArrayType, StringType

TRAIT_0 = 0
TRAIT_1 = 1
TRAIT_2 = 2


def flag_to_list(flag: int) -> List[str]:
    trait_list = []
    if flag & (1 << TRAIT_0):
        trait_list.append("TRAIT_0")
    elif flag & (1 << TRAIT_1):
        trait_list.append("TRAIT_1")
    elif flag & (1 << TRAIT_2):
        trait_list.append("TRAIT_2")
    return trait_list


flag_to_list_udf = udf(lambda x: None if x is None else flag_to_list(x),
                       ArrayType(StringType()))

# Create dummy data to test
data = [
    { "name": "John Doe", "traits": 5 },
    { "name": "Jane Doe", "traits": 2 },
    { "name": "Jane Roe", "traits": 0 },
    { "name": "John Roe", "traits": 6 },
]
df = spark.createDataFrame(data, 'name STRING, traits INT')
df.show()
# +--------+------+
# |    name|traits|
# +--------+------+
# |John Doe|     5|
# |Jane Doe|     2|
# |Jane Roe|     0|
# |John Roe|     6|
# +--------+------+

df = df.withColumn('traits_processed', flag_to_list_udf(df['traits']))
df.show()
# +--------+------+----------------+
# |    name|traits|traits_processed|
# +--------+------+----------------+
# |John Doe|     5|       [TRAIT_0]|
# |Jane Doe|     2|       [TRAIT_1]|
# |Jane Roe|     0|              []|
# |John Roe|     6|       [TRAIT_1]|
# +--------+------+----------------+

如果您不想创建新列，可以将traits_processed 替换为traits。