python之如何将多列作为参数传递给 pyspark write repartition()

我创建了一个函数，它使用一些参数将数据帧写出到 s3 位置。除重新分区列表参数外的所有工作。它因错误而失败:raise TypeError("numPartitions should be an int or Column")

年份是int型列，日期是sate型列。当我将 2 列硬编码为 .repartition("parti_create_date", "parti_hour") 它有效。 我尝试将它们作为列表、字符串和列提供。 似乎没有任何效果。

parti_list = ["parti_year", "parti_create_date", "parti_hour"] 
re_parti_list = ["parti_create_date", "parti_hour"] 
 
def spark_write(in_df, write_tgt_loc, parti_list, re_parti_list, tgt_file_format, write_mode, tgt_file_compression): 
(in_df 
            .repartition(re_parti_list)  #(re_parti_str) 
            .write 
            .partitionBy(parti_str) 
            .mode(write_mode).format(tgt_file_format) 
            .option('compression', tgt_file_compression).option("nullValue", "null").option("treatEmptyValuesAsNulls,", "true") 
            .save(write_tgt_loc)) 
 
spark_write(tgt_df, "s3://bucket/out/", parti_list, re_parti_list, "parquet", "overwrite","snappy") 

你能帮我弄清楚如何在 PySpark 中将重新分区的列作为参数传递吗？

请您参考如下方法:

重新分区期望either int or column , 所以我们需要通过col("<col_name>")到数据框。

Example:

df=spark.createDataFrame([(1,'a',),(2,'b',),(3,'c',)],['id','name']) 
df.rdd.getNumPartitions() #number of partitions in df 
1 

Repartition on int:

df.repartition(10).rdd.getNumPartitions() #repartition to 10  
 
10 

Repartition on columns:

df.repartition(col("id"),col("name")).rdd.getNumPartitions() #repartition on columns 
 
200 

Dynamic repartition on columns:

df.repartition(*[col(c) for c in df.columns]).rdd.getNumPartitions() 
 
200 

映射您的 columns list至 column键入而不是 string 然后传递 repartition. 中的列名

For your case try this way:

df.repartition(*[col(c) for c in re_parti_list]) 
            .write 
            .partitionBy(parti_str) 
            .mode(write_mode).format(tgt_file_format) 
            .option('compression', tgt_file_compression).option("nullValue", "null").option("treatEmptyValuesAsNulls,", "true") 
            .save(write_tgt_loc)) 

In scala:

df.repartition(df.columns.map(c => col(c)):_*).rdd.getNumPartitions 
200