PySpark -添加一个递增的数字列，该列根据另一个列值的变化重置为1

from pyspark.sql import SparkSession
from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

# 创建SparkSession
spark = SparkSession.builder.getOrCreate()

# 创建示例数据
data = [("A",), ("A",), ("B",), ("B",), ("B",), ("C",)]

# 创建DataFrame
df = spark.createDataFrame(data, ["col1"])

# 定义窗口规范
windowSpec = Window.partitionBy("col1").orderBy("col1")

# 添加递增的数字列
df = df.withColumn("col2", row_number().over(windowSpec))

# 显示结果
df.show()

+----+----+
|col1|col2|
+----+----+
|   A|   1|
|   A|   2|
|   B|   1|
|   B|   2|
|   B|   3|
|   C|   1|
+----+----+