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

Question

我创建了一个 function ，它使用一些参数将 dataframe 写入 s3 位置。 除重新分区列表参数外的所有工作。 它失败并出现错误： raise TypeError("numPartitions should be an int or Column")

年份是一个 int 列，日期是一个状态类型列。 当我将 2 列硬编码为.repartition("parti_create_date", "parti_hour")时，它可以工作。 我尝试以列表、字符串和 col 的形式提供它们。 似乎没有任何效果。

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 中的参数传递吗？

Answer 1

重新分区需要either int or column ，因此我们需要将col("<col_name>")传递给 dataframe。

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

map 您的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