Pyspark: function 加入可變列數

Question

我想知道是否有辦法自動化...我想做一個 function 我會告訴我要加入多少列。 如果我有 3 列的 dataFrame 並給出參數“number_of_columns = 3”，那么它將加入列：0、1、2。但如果我有 7 列的 dataFrame 並給出參數“number_of_columns = 7”，那么它將連接列：0、1、2、3、4、5、6。列的名稱始終相同：從“0”到“number_of_columns-1”。

有沒有辦法做到這一點？ 或者如果我要合並其他數量的列，我必須有另一個 function？

def my_function(spark_column, name_of_column):
    new_spark_column = spark_column.withColumn(name_of_column, concat_ws("", 
                                                   col("0").cast("Integer"), 
                                                   col("1").cast("Integer"),
                                                   col("2").cast("Integer"),
                                                   col("3").cast("Integer"),
                                                   col("4").cast("Integer"),
                                                   col("5").cast("Integer"),
                                                   col("6").cast("Integer") ))

Answer 1

您可以使用列表推導來執行此操作：

from pyspark.sql.functions import concat_ws, col

def my_function(spark_column, n_cols, name_of_column):
    new_spark_column = spark_column.withColumn(
        name_of_column, 
        concat_ws("", *[col(c).cast("Integer") for c in spark_column.columns[:n_cols]])
    )
    return new_spark_column