Разделение файлов по размеру в PySpark

Question

Я использую PySpark для разбиения df и записи 10 разделов на S3:

df.repartition(10).write.format("json").mode("OverWrite").save("s3Bucket")

Как я могу разделить df, но основываясь на размере разделов, а не на количестве разделов. Пример: я хочу иметь n разделов по 50 МБ.

Answer 1

Если DataFrame загружен из файлов, расположенных в вашем контейнере, вы можете получить размер входных файлов и использовать его для вычисления количества разделов.

Примерно так должно работать:

data_path = "/s3Bucket/path/"
PARTITION_BLOCK_SIZE = 50000000 # 50MB

# using Hadoop FS API to get data len
Path = sc._gateway.jvm.org.apache.hadoop.fs.Path

fs = Path(data_path).getFileSystem(sc._jsc.hadoopConfiguration())
data_size = fs.getFileStatus(Path(data_path)).getLen()

# repartition if size > 50MB
if (data_size / PARTITION_BLOCK_SIZE) > 1:
    nb_partitions = math.ceil(data_size / PARTITION_BLOCK_SIZE)
    df = df.repartition(nb_partitions)