How to pass multi columns as parameters to pyspark write repartition()

user_5191381 0
Viewed 1205

I have created a function that writes out a dataframe to an s3 location with some parameters. All works except for re-partition list parameter. It fails with the error: raise TypeError("numPartitions should be an int or Column")

The year is an int column and the date is a sate type column. When I hard-code the 2 columns as .repartition("parti_create_date", "parti_hour") it works. I tried supplying them as a list, string, and col. Nothing seems to work.

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

Can you please help me figure out how to pass the re-partition columns as parameters in PySpark?

pyspark apache-spark python
1 Answers

Repartition expects either int or column, So we need to pass col("<col_name>") to the 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 your columns list to column type instead of string then pass the column names in 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
user_76326950
Related