如何基于PySpark DataFrame首两行模式批量更新列值？

针对大数据集场景，我们可以利用PySpark的窗口函数和取模运算实现需求，全程分布式处理，无需循环或列表操作，避免将大量数据拉到Driver端。

方法一：纯分布式窗口函数实现（无需收集数据到Driver） #

这种方式全程在Executor端处理，完全不需要将任何数据拉取到Driver，适合极端大数据场景：

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

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

# 示例初始DataFrame
data = [(1, 2.54), (2, 3.21), (3, 5.67), (4, 8.90), (5, 1.23), (6, 4.56)]
df = spark.createDataFrame(data, ["Col1", "Col2"])

# 定义全局排序窗口（按Col1排序保证行序）
global_window = Window.orderBy("Col1")

# 提取首两行的Col2值作为固定模式
df_with_pattern = df.withColumn("pattern_1", F.first("Col2").over(global_window)) \
                    .withColumn("pattern_2", F.lag("Col2", 1).over(global_window)) \
                    .withColumn("pattern_2", F.first("pattern_2").over(global_window))

# 添加行号并根据取模结果循环赋值
result_df = df_with_pattern.withColumn("row_num", F.row_number().over(global_window)) \
                           .withColumn("new_Col2",
                                       F.when(F.col("row_num") % 2 == 1, F.col("pattern_1"))
                                        .otherwise(F.col("pattern_2"))) \
                           .drop("pattern_1", "pattern_2", "row_num")

# 查看结果
result_df.show()

方法二：轻量收集首行值后分布式处理（更简洁） #

由于仅需收集前两行的两个值，数据量极小，不会对Driver造成压力，代码更简洁：

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

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

# 示例初始DataFrame
data = [(1, 2.54), (2, 3.21), (3, 5.67), (4, 8.90), (5, 1.23), (6, 4.56)]
df = spark.createDataFrame(data, ["Col1", "Col2"])

# 提取首两行的Col2值作为模式（仅两个值，无性能问题）
pattern_vals = df.limit(2).select("Col2").rdd.flatMap(lambda x: x).collect()
val1, val2 = pattern_vals[0], pattern_vals[1]

# 定义全局窗口添加行号，按取模结果循环赋值
global_window = Window.orderBy("Col1")
result_df = df.withColumn("row_num", F.row_number().over(global_window)) \
              .withColumn("new_Col2",
                          F.when(F.col("row_num") % 2 == 1, val1)
                           .otherwise(val2)) \
              .drop("row_num")

# 查看结果
result_df.show()

核心逻辑说明 #

1. 利用row_number()生成连续行号，保证行序与初始DataFrame一致；
2. 通过行号对2取模，判断当前行属于模式中的第几个值（余数1取第一个值，余数0取第二个值）；
3. 两种方法均采用分布式算子处理，不会触发全量数据的Shuffle或Driver端内存压力，完全适配大数据集。

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