pyspark function.lag on condition

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import lag, col, when

# 初始化 Spark 会话
spark = SparkSession.builder.appName("LagExample").getOrCreate()

# 假设我们有一个 DataFrame，包含用户ID和登录日期
data = [("user1", "2023-01-01"), ("user1", "2023-01-02"), ("user1", "2023-01-04"),
        ("user2", "2023-01-01"), ("user2", "2023-01-03")]
columns = ["userId", "loginDate"]

df = spark.createDataFrame(data, columns)

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

# 使用 lag 函数获取前一行的登录日期
df_with_lag = df.withColumn("prevLoginDate", lag("loginDate").over(windowSpec))

# 计算连续登录天数
df_with_consecutive_days = df_with_lag.withColumn(
    "consecutiveDays",
    when(col("prevLoginDate") == col("loginDate") - 1, 1).otherwise(0)
)

df_with_consecutive_days.show()