PySpark：根据指定元素的首尾出现截断列中的列表

你可以通过自定义UDF（用户定义函数）来实现这个需求，核心逻辑是找到列表中首次和末次出现目标元素的位置，然后对列表进行切片。以下是完整的实现代码：

步骤1：定义输入DataFrame并指定列名 #

首先给输入DataFrame明确列名，方便后续操作：

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

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

df_in = spark.createDataFrame(
    [(1,['A', 'B', 'A', 'F', 'C', 'D']),
     (2,['F', 'C', 'B', 'X', 'A', 'D']),
     (3,['L', 'A', 'B', 'M', 'C'])],
    schema=["id", "items"]
)

步骤2：编写过滤逻辑函数 #

定义一个函数，接收列表和两个目标元素，返回符合要求的子列表：

def filter_between_first_last(items, e1, e2):
    # 收集所有出现e1或e2的元素索引
    target_indices = [idx for idx, item in enumerate(items) if item in (e1, e2)]
    if not target_indices:
        # 若列表中无目标元素，返回原列表（可根据需求修改为返回空列表）
        return items
    # 取首次和末次出现的索引
    first_pos = target_indices[0]
    last_pos = target_indices[-1]
    # 切片获取目标区间的元素（左闭右闭）
    return items[first_pos:last_pos + 1]

步骤3：注册UDF并应用到DataFrame #

将函数注册为PySpark可识别的UDF，然后对目标列进行转换：

# 注册UDF，指定返回类型为字符串数组
filter_list_udf = udf(filter_between_first_last, ArrayType(StringType()))

# 生成结果DataFrame
df_out = df_in.select(
    "id",
    filter_list_udf("items", lit("A"), lit("C")).alias("items")
)

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

运行上述代码后，输出结果将与你期望的df_out完全一致：

+---+---------------------+
|id |items                |
+---+---------------------+
|1  |[A, B, A, F, C]      |
|2  |[C, B, X, A]         |
|3  |[A, B, M, C]         |
+---+---------------------+

补充说明 #

• 如果需要处理其他目标元素，只需修改lit("A")和lit("C")为对应的元素即可。
• 若列表中不存在任何目标元素，当前逻辑会返回原列表，你可以根据实际需求修改函数中的返回值（例如改为return []）。

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