Pyspark选择最接近date2值的date1列的值，如果小于date1，则优先选择date2，分组ID

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

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

data = [(1, "2022-01-01", "2022-01-02"),
        (1, "2022-01-03", "2022-01-04"),
        (2, "2022-01-05", "2022-01-06"),
        (2, "2022-01-07", "2022-01-08")]
df = spark.createDataFrame(data, ["ID", "date1", "date2"])

df = df.withColumn("date1", col("date1").cast("date"))
df = df.withColumn("date2", col("date2").cast("date"))

windowSpec = Window.partitionBy("ID").orderBy(expr("abs(datediff(date1, date2))"))
df = df.withColumn("closest_date1", when(col("date1") < col("date2"), col("date2")).otherwise(col("date1")))
df = df.withColumn("row_number", row_number().over(windowSpec))
df = df.filter(col("row_number") == 1).drop("row_number")