如何编写可以接受参数数量可变的参数的 Pyspark 函数？

Question

我编写了一个函数，我想对其进行修改以使其具有可以采用一个或多个参数的参数，但我无法使其正常工作。

def get_recent_date(input_df, *partion_col, order_col):
  w = Window().partitionBy(partition_col)\
  .orderBy(desc(order_col))
  output_df= input_df.withColumn('DenseRank', dense_rank().over(w))
  return output_df

我希望该函数运行，以便partition_col 可以采用可变数量的参数。在下面的示例 1 中，partition_col= 'event_category'，在示例 2 中，partition_col = 'event_category' 和 'participant_category'。我尝试过以多种方式运行它，但经常收到错误“TypeError：只能将 str （不是“元组”）连接到 str”。预先感谢您的帮助！

ex 1: get_recent_date(input, 'event_category', 'event_date')

ex 2: get_recent_date(input, 'event_category', 'participant_category', 'event_date')

Answer 1

*partion_col

允许您传递可变数量的非关键字参数，其中

partion_col

将表示函数内的非关键字参数

tuple

。因此，您需要解压列名称元组，以便

pyspark

正确使用可变长度分区。

更换

w = Window().partitionBy(partition_col)\

与

w = Window().partitionBy(*partition_col)\

你应该可以走了。

可复制示例：

from pyspark.sql import SparkSession, Window
from pyspark.sql.functions import desc, dense_rank

spark = SparkSession.builder.appName('spark_session').getOrCreate()

data = [
    (100, 1, 2, 1),
    (100, 1, 1, -1),
    (200, 1, 3, 1),
    (200, 1, 3, 4)   
]

df = spark.createDataFrame(data, ("col_1", "col_2", "col_3", 'order_col'))

df.show()

# +-----+-----+-----+---------+
# |col_1|col_2|col_3|order_col|
# +-----+-----+-----+---------+
# |  100|    1|    2|        1|
# |  100|    1|    1|       -1|
# |  200|    1|    3|        1|
# |  200|    1|    3|        4|
# +-----+-----+-----+---------+

def get_recent_date(input_df, *partition_col, order_col):
    w = Window().partitionBy(*partition_col)\
    .orderBy(desc(order_col))
    output_df= input_df.withColumn('DenseRank', dense_rank().over(w))
    return output_df


new_df = get_recent_date(
    df, 'col_2', order_col='order_col'
)

new_df.show()

# +-----+-----+-----+---------+---------+
# |col_1|col_2|col_3|order_col|DenseRank|
# +-----+-----+-----+---------+---------+
# |  200|    1|    3|        4|        1|
# |  100|    1|    2|        1|        2|
# |  200|    1|    3|        1|        2|
# |  100|    1|    1|       -1|        3|
# +-----+-----+-----+---------+---------+

new_df = get_recent_date(
    df, 'col_2', 'col_1', order_col='order_col'
)

new_df.show()

# |col_1|col_2|col_3|order_col|DenseRank|
# +-----+-----+-----+---------+---------+
# |  100|    1|    2|        1|        1|
# |  100|    1|    1|       -1|        2|
# |  200|    1|    3|        4|        1|
# |  200|    1|    3|        1|        2|
# +-----+-----+-----+---------+---------+