Проблема в том что нужно контролировать размер файлов а не их количество
Помимо того, что я советовал в соседней конфе (москау спарк) есть ещё вариант для отважных: свой DataSource, в котором можно сразу делать нужные партиции, Есть DataSourceV1 - старый, но надёжный и V2 - новый, получше продуман, но может изменяться с минорными версиями (Spark 2 интерфейсы V2 были evolving), Spark 3 не смотрел ещё)