在合并列函数中应用 *args 并指示字符串值

Posted 2023-03-23

【中文标题】在合并列函数中应用 *args 并指示字符串值

【英文标题】：Applying *args and indicating string values in merging columns function

【发布时间】：2022-01-22 04:06:36

【问题描述】：

我正在使用以下函数，但想应用args（而不是column1、column2）来使列的数量可变。我还想给出参数应该是字符串值的参数。

有人知道如何实现吗？

def create_new_column_by_merging_columns(spark_df, new_column, column_1, column_2):
    spark_df_with_new_column = spark_df.withColumn(new_column, sf.concat(sf.col(column_1), sf.lit('_'), sf.col(column_2)))
    
    return spark_df_with_new_column


new_df_with_merged_columns = create_new_column_by_merging_columns(spark_df, new_column, column_1, column_2)

【参考方案1】：

您可以简单地使用 concat_ws 函数，它完全符合您的要求：

pyspark.sql.functions.concat_ws(sep, *cols)

将多个输入字符串列连接成一个 字符串列，使用给定的分隔符。

import pyspark.sql.functions as F

df = df.withColumn(new_column, F.concat_ws("_", column1, column2, ...))

但是，如果你真的需要一个函数，那么你可以定义如下：

def merge_columns(df, new_column, *cols):
    return df.withColumn(new_column, F.concat_ws("_", *cols))