如何拆分PySpark函数并将窗口函数传递至另一函数？

我有一个可正常运行的PySpark函数，希望将其拆分为两个函数。尝试将Spark窗口函数作为参数传递给另一个函数，约束条件是函数仅接收列名参数，不能传递DataFrame。以下是我的DataFrame定义、原可运行函数以及拆分后无法运行的两个函数，请求帮助解决：

DataFrame定义 #

df2 = df.select(*cols, F.to_date("date").alias('datenew'))

原可运行函数 #

def price_col(df2, id_col, subcatid_col, start_date_col, year_col, statusid_col, price_col):
    window_spec_st = Window().partitionBy(F.col(id_col), F.col(subcatid_col)).orderBy(F.col(start_date_col))
    return df2.where(
               (F.col(price_col) > F.lit(0)) & 
               (F.col(statusid_col) == F.lit(1)))
        .withColumn("row_number", F.row_number().over(window_spec_st))
        .where((F.col("row_number") == F.lit(1)))
        .drop("row_number", "end_date", "end__date", start_date_col)

拆分后无法运行的函数 #

def set_windows_spec(id_col, subcatid_col, date_col, desc_order):
    order_by_col = F.col(date_col).desc() if desc_order else F.col(date_col).asc()
    return f'Window().partitionBy(F.col(id_col), F.col(subcatid_col)).orderBy(order_by_col)'


def conditions(price_col, statusid_col,id_col, subcatid_col, date_col,desc_order, set_windows_spec):
    return (F.when(
        (F.col(price_col) > F.lit(0)) &
        (F.col(statusid_col) == F.lit(1)),
         F.col(price_col))
            .withColumn("row_number", F.row_number().over(set_windows_spec(id_col, subcatid_col, date_col, desc_order)))
            .where(F.col("row_number") == F.lit(1))
            .alias("new_price_col"))

问题原因分析 #

1. 窗口函数返回格式错误：set_windows_spec函数返回的是字符串形式的代码，而非实际的Window对象。F.row_number().over()需要接收真实的窗口规范对象，字符串无法被识别。
2. Column与DataFrame操作混淆：F.when()返回的是Column对象，而withColumn、where是DataFrame的方法，不能直接在Column对象上调用，这是核心错误。

修正后的代码 #

1. 生成窗口规范的函数 #

返回真实的Window对象，而非字符串：

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

def get_window_spec(id_col, subcatid_col, date_col, desc_order=False):
    order_expr = F.col(date_col).desc() if desc_order else F.col(date_col).asc()
    return Window.partitionBy(F.col(id_col), F.col(subcatid_col)).orderBy(order_expr)

2. 生成列操作逻辑的函数 #

仅生成列表达式和过滤条件，符合“仅接收列参数”的约束：

def get_price_processing_cols(price_col, statusid_col, window_spec):
    # 基础过滤条件
    base_filter = (F.col(price_col) > F.lit(0)) & (F.col(statusid_col) == F.lit(1))
    # 窗口行号列
    row_num_col = F.row_number().over(window_spec).alias("row_number")
    # 最终的目标价格列（仅保留符合条件且为分组第一行的价格）
    target_price_col = F.when(base_filter & (row_num_col == 1), F.col(price_col)).alias("new_price_col")
    return base_filter, row_num_col, target_price_col

使用示例 #

在DataFrame层面完成后续操作（符合约束，函数仅处理列相关逻辑）：

# 获取窗口规范
window_spec = get_window_spec("id_col", "subcatid_col", "start_date_col")
# 获取过滤条件、行号列、目标价格列
base_filter, row_num_col, target_price_col = get_price_processing_cols("price_col", "statusid_col", window_spec)

# 应用到DataFrame
result_df = df2.filter(base_filter)\
               .withColumn("row_number", row_num_col)\
               .filter(F.col("row_number") == 1)\
               .drop("row_number", "end_date", "end__date", "start_date_col")\
               .withColumn("new_price_col", target_price_col)

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