pyspark: Efficiently have partitionBy write to same number of total partitions as original table

by

You’ve got several options. In my code below I’ll assume you want to write in parquet, but of course you can change that.

(1) df.repartition(numPartitions, *cols).write.partitionBy(*cols).parquet(writePath)

This will first use hash-based partitioning to ensure that a limited number of values from COL make their way into each partition. Depending on the value you choose for numPartitions, some partitions may be empty while others may be crowded with values — for anyone not sure why, read this. Then, when you call partitionBy on the DataFrameWriter, each unique value in each partition will be placed in its own individual file.

Warning: this approach can lead to lopsided partition sizes and lopsided task execution times. This happens when values in your column are associated with many rows (e.g., a city column — the file for New York City might have lots of rows), whereas other values are less numerous (e.g., values for small towns).

(2) df.sort(sortCols).write.parquet(writePath)

This options works great when you want (1) the files you write to be of nearly equal sizes (2) exact control over the number of files written. This approach first globally sorts your data and then finds splits that break up the data into k evenly-sized partitions, where k is specified in the spark config spark.sql.shuffle.partitions. This means that all values with the same values of your sort key are adjacent to each other, but sometimes they’ll span a split, and be in different files. This, if your use-case requires all rows with the same key to be in the same partition, then don’t use this approach.

There are two extra bonuses: (1) by sorting your data its size on disk can often be reduced (e.g., sorting all events by user_id and then by time will lead to lots of repetition in column values, which aids compression) and (2) if you write to a file format the supports it (like Parquet) then subsequent readers can read data in optimally by using predicate push-down, because the parquet writer will write the MAX and MIN values of each column in the metadata, allowing the reader to skip rows if the query specifies values outside of the partition’s (min, max) range.

Note that sorting in Spark is more expensive than just repartitioning and requires an extra stage. Behind the scenes Spark will first determine the splits in one stage, and then shuffle the data into those splits in another stage.

(3) df.rdd.partitionBy(customPartitioner).toDF().write.parquet(writePath)

If you’re using spark on Scala, then you can write a customer partitioner, which can get over the annoying gotchas of the hash-based partitioner. Not an option in pySpark, unfortunately. If you really want to write a custom partitioner in pySpark, I’ve found this is possible, albeit a bit awkward, by using rdd.repartitionAndSortWithinPartitions:

df.rdd \
  .keyBy(sort_key_function) \  # Convert to key-value pairs
  .repartitionAndSortWithinPartitions(numPartitions=N_WRITE_PARTITIONS, 
                                      partitionFunc=part_func) \
  .values() # get rid of keys \
.toDF().write.parquet(writePath)

Maybe someone else knows an easier way to use a custom partitioner on a dataframe in pyspark?

More Related Contents:

  1. How to melt Spark DataFrame?
  2. Unpivot in spark-sql/pyspark
  3. Spark Dataframe validating column names for parquet writes
  4. How to check if spark dataframe is empty?
  5. How to improve performance for slow Spark jobs using DataFrame and JDBC connection?
  6. TypeError: Column is not iterable – How to iterate over ArrayType()?
  7. How to save/insert each DStream into a permanent table
  8. How to fix ‘TypeError: an integer is required (got type bytes)’ error when trying to run pyspark after installing spark 2.4.4
  9. Pyspark : forward fill with last observation for a DataFrame
  10. Adding a group count column to a PySpark dataframe
  11. PySpark: how to resample frequencies
  12. Multiple Spark applications with HiveContext
  13. How to loop through each row of dataFrame in pyspark
  14. Explode array data into rows in spark [duplicate]
  15. Why does Spark think this is a cross / Cartesian join
  16. How to find count of Null and Nan values for each column in a PySpark dataframe efficiently?
  17. Pyspark: Pass multiple columns in UDF
  18. Rename more than one column using withColumnRenamed
  19. Spark gives a StackOverflowError when training using ALS
  20. AWS EMR – ModuleNotFoundError: No module named ‘pyarrow’
  21. Not able to cat dbfs file in databricks community edition cluster. FileNotFoundError: [Errno 2] No such file or directory:
  22. PySpark: How to fillna values in dataframe for specific columns?
  23. pyspark: rolling average using timeseries data
  24. How to exclude multiple columns in Spark dataframe in Python
  25. Stratified sampling with pyspark
  26. Get CSV to Spark dataframe
  27. How to turn off scientific notation in pyspark?
  28. Spark ALS predictAll returns empty
  29. Why is the fold action necessary in Spark?
  30. How to add a SparkListener from pySpark in Python?

Leave a Comment