pyspark:连接可变列数的函数

jchrr9hc  于 2021-07-09  发布在  Spark
关注(0)|答案(1)|浏览(243)

我想知道是否有一种方法可以自动化。。。我想做一个函数,在其中我会告诉,我想加入多少列。如果我有一个有3列的dataframe,并给出一个参数“numberofcolumns=3”,那么它将连接列:0、1、2。但是如果我有一个包含7列的dataframe,并给出一个参数“numberofcolumns=7”,那么它将连接列:0、1、2、3、4、5、6。列的名称总是相同的:从“0”到“number\ of\ columns-1”。
有什么办法吗?或者我必须有另一个函数,如果我有其他数量的列合并?

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") ))
mfuanj7w

mfuanj7w1#

您可以使用列表理解来执行此操作:

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

相关问题