Package: repartipy
RepartiPyRepartiPy helps you to elaborately handle PySpark DataFrame partition size.Possible Use CasesRepartition your DataFrame precisely,without knowing the whole DataFrame size(i.e.Dynamic Repartition)Estimate your DataFrame sizewith more accuracyWhy RepartiPyAlthough SparkSizeEstimatorcan be used to estimate a DataFrame size, it is not accurate sometimes.
RepartiPy usesSpark's execution plan statisticsin order to provide a roundabout way.
It suggests two approaches to achieve this:reaprtipy.SizeEstimatorreaprtipy.SamplingSizeEstimatorreaprtipy.SizeEstimatorRecommended when your executor resource (memory) is affordable to cache the whole DataFrame.SizeEstimatorjust simply caches the whole Dataframe into the memory and extract the execution plan statistics.repartipy.SamplingSizeEstimatorRecommended when your executor resource (memory) isNOTaffordable to cache the whole dataframe.SamplingSizeEstimatoruses 'disk write and re-read (HDFS)' approach behind the scene for two reasons:Prevent double read from the source data (check point) -> better performanceReduce partition skewness by reading again (leverage MaxSplitBytes) -> better sampling resultTherefore,you must have HDFS settings on your cluster and enough disk space.This may not be accurate compared toSizeEstimatordue to sampling.
If you want more accurate results, tune thesample_countoption properly.
Additionally, this approach will be slower thanSizeEstimatorasSamplingSizeEstimatorrequires disk I/O and additional logics.How To UseSetuppipinstallrepartipyPrerequisitefrompyspark.sqlimportSparkSessionspark=SparkSession.builder.getOrCreate()input_data=[(1,"Seoul"),(2,"Busan"),]df=spark.createDataFrame(data=input_data,schema=["id","location"])get_desired_partition_count()Calculate ideal number of partitions for a DataFrameSizeEstimator will suggestdesired_partition_count, so that each partition can havedesired_partition_size_in_bytes(default: 1GiB) after repartition.reproduce()produces exactly the samedf, but internally reproduced by SizeEstimator for better performance.SizeEstimatorreproducesdffromMemory(Cache).SamplingSizeEstimatorreproducesdffromDisk(HDFS).with SizeEstimatorimportrepartipyone_gib_in_bytes=1073741824withrepartipy.SizeEstimator(spark=spark,df=df)asse:desired_partition_count=se.get_desired_partition_count(desired_partition_size_in_bytes=one_gib_in_bytes)se.reproduce().repartition(desired_partition_count).write.save("your/write/path")# orse.reproduce().coalesce(desired_partition_count).write.save("your/write/path")with SamplingSizeEstimatorimportrepartipyone_gib_in_bytes=1073741824withrepartipy.SamplingSizeEstimator(spark=spark,df=df,sample_count=10)asse:desired_partition_count=se.get_desired_partition_count(desired_partition_size_in_bytes=one_gib_in_bytes)se.reproduce().repartition(desired_partition_count).write.save("your/write/path")# orse.reproduce().coalesce(desired_partition_count).write.save("your/write/path")estimate()Estimate size of a DataFramewith SizeEstimatorimportrepartipywithrepartipy.SizeEstimator(spark=spark,df=df)asse:df_size_in_bytes=se.estimate()with SamplingSizeEstimatorimportrepartipywithrepartipy.SamplingSizeEstimator(spark=spark,df=df,sample_count=10)asse:df_size_in_bytes=se.estimate()BenchmarkOverall, there appears to be a slight performance loss when employing RepartiPy.
This benchmark compares therunning time of spark jobsin the following two cases to give a rough estimate:StaticRepartition (repartition without RepartiPy)# e.g.df.repartition(123).write.save("your/write/path")DynamicRepartition (repartition with RepartiPy)# e.g.withrepartipy.SizeEstimator(spark=spark,df=df)asse:desired_partition_count=se.get_desired_partition_count(desired_partition_size_in_bytes=one_gib_in_bytes)se.reproduce().repartition(desired_partition_count).write.save("your/write/path")All the other conditions remain the sameexcept the usage of RepartiPy.NoteBenchmark results provided are for brief reference only, not absolute.
Actual performance metrics can vary depending on your own circumstances (e.g. your data, your spark code, your cluster resources, ...).SizeEstimatorDataFrame Size ~= 256 MiB (decompressed size)StaticDynamicRunning Time8.5 min8.6 minSamplingSizeEstimatorDataFrame Size ~= 241 GiB (decompressed size)StaticDynamicRunning Time14 min16 min