PySpark DataFrame：将布尔字符串列表转换为可变长嵌套列表

需求说明 #

将DataFrame中sequence列的每个布尔字符串（如'1000'），映射为预定义orientation列表中对应位置的字符组成的子列表，保持原序列的顺序。其中orientation = ['L', 'D', 'U', 'R']，布尔字符串的每一位'1'对应orientation中相同索引的字符。

方案一：使用PySpark内置函数（推荐，性能更优） #

适用于Spark 3.0及以上版本，利用原生内置函数实现，无需自定义UDF，性能更高效。

代码实现 #

from pyspark.sql import SparkSession, Row
from pyspark.sql import functions as F

# 初始化SparkSession
spark = SparkSession.builder.appName("OrientationTransform").getOrCreate()

# 构造原始DataFrame
df = spark.sparkContext.parallelize([
    Row(sequence=['1000', '0011']),
    Row(sequence=['0100', '1001', '0010']),
    Row(sequence=['1011', '0100']),
    Row(sequence=['1111'])
]).toDF()

# 定义方向标识列表
orientation = ['L', 'D', 'U', 'R']

# 执行转换逻辑
transformed_df = df.withColumn(
    "sequence",
    F.transform(
        "sequence",
        lambda s: F.filter(
            F.zip_with(
                F.split(s, ""),  # 将单个布尔字符串拆分为字符数组
                F.array(*[F.lit(o) for o in orientation]),  # 将orientation转为Spark数组
                lambda char, orient: F.when(char == '1', orient).otherwise(None)
            ),
            lambda x: x.isNotNull()  # 过滤掉对应'0'的空值
        )
    )
)

# 查看结果
transformed_df.show(truncate=False)

逻辑说明 #

1. 使用transform遍历sequence数组中的每个布尔字符串；
2. 用split将字符串拆分为单个字符的数组；
3. 通过zip_with将字符数组与orientation数组配对，仅保留字符为'1'对应的方向标识；
4. 最后用filter移除空值，得到目标子列表。

方案二：使用自定义UDF #

若使用低版本Spark（不支持zip_with等函数），可通过自定义UDF实现：

代码实现 #

from pyspark.sql import SparkSession, Row
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType

spark = SparkSession.builder.appName("OrientationTransformUDF").getOrCreate()

# 构造原始DataFrame
df = spark.sparkContext.parallelize([
    Row(sequence=['1000', '0011']),
    Row(sequence=['0100', '1001', '0010']),
    Row(sequence=['1011', '0100']),
    Row(sequence=['1111'])
]).toDF()

orientation = ['L', 'D', 'U', 'R']

# 定义转换函数
def map_orientation(seq_list):
    result = []
    for s in seq_list:
        # 遍历字符串的每个字符，收集对应方向标识
        sub_list = [orientation[i] for i, char in enumerate(s) if char == '1']
        result.append(sub_list)
    return result

# 注册UDF并指定返回类型
orientation_udf = udf(map_orientation, ArrayType(ArrayType(StringType())))

# 应用UDF完成转换
transformed_df = df.withColumn("sequence", orientation_udf("sequence"))

transformed_df.show(truncate=False)

逻辑说明 #

自定义函数遍历每个布尔字符串，根据字符'1'的索引提取orientation中对应的元素，组成子列表后返回。

内容的提问来源于stack exchange，提问作者Wilson