如何在PySpark中实现类似grep -B的匹配行及上一行输出功能

当然可以实现！PySpark完全能达成你想要的——筛选出包含指定关键词的行及其上一行，效果类似grep -B 1。下面我就用你提供的示例数据，一步步演示具体实现方案，还会考虑边界情况哦。

步骤1：准备环境与示例数据 #

首先我们需要初始化SparkSession，并创建你的示例DataFrame。注意：Spark是分布式计算框架，默认不保证数据的顺序，所以我们需要先从字符串里提取日期字段，用来后续排序确保行的顺序正确。

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

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

# 你的示例数据
sample_data = [
    ("2021-08-30 end active",),
    ("2021-09-01 end inactive",),
    ("2021-09-02 start active",),
    ("2021-09-03 end active",)
]

# 创建基础DataFrame
df = spark.createDataFrame(sample_data, ["string"])

# 提取日期列（用于后续排序，保证行顺序正确）
df = df.withColumn("date", F.to_date(F.split(F.col("string"), " ")[0], "yyyy-MM-dd"))

步骤2：为每行添加有序行号 #

我们通过窗口函数row_number()，按日期升序为每行分配唯一行号，这样就能精准定位匹配行的上一行。

# 定义窗口：按日期升序排序，确保行号顺序和数据时序一致
sort_window = Window.orderBy("date")

# 添加行号列
df_with_row_num = df.withColumn("row_num", F.row_number().over(sort_window))

步骤3：定位目标行并筛选 #

首先找出所有包含关键词"start"的行的行号，然后生成需要保留的行号列表（匹配行号 + 匹配行号减1），最后筛选出这些行即可。同时我们会处理边界情况：如果匹配行是第一行，就跳过不存在的"上一行"。

# 收集所有包含"start"的行的行号
matching_rows = df_with_row_num.filter(F.col("string").contains("start")).select("row_num")
matching_row_nums = matching_rows.rdd.flatMap(lambda x: x).collect()

# 生成目标行号：匹配行号 + 匹配行的上一行行号（排除行号<=0的情况）
target_row_nums = matching_row_nums + [num - 1 for num in matching_row_nums if num > 1]

# 筛选目标行，并移除辅助列（date和row_num）
result_df = df_with_row_num.filter(F.col("row_num").isin(target_row_nums)).drop("date", "row_num")

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

运行结果 #

执行上述代码后，你会得到期望的输出：

+------------------------+
|string                  |
+------------------------+
|2021-09-01 end inactive |
|2021-09-02 start active |
+------------------------+

关键注意点 #

• 数据顺序的重要性：必须明确指定排序字段（比如示例中的日期），否则Spark生成的行号会混乱，导致无法正确匹配上一行。
• 边界处理：当匹配行是第一行时，num - 1会得到0，这时候我们通过if num > 1过滤掉，避免筛选出不存在的行。

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