Ток там используют нормальный сэмплинг
Посмотрел, используется Reservoir sampling. В принципе, в percentile_approx тоже приблизительный расчёт с эффективным использованием памяти и в первой прототипной реализации всё было хорошо, проблемы появились, когда потребовалось сделать распределение динамическим. Метод sample датасета никогда проблем не вызывал, он активно используется и качество выборки на миллиардах записей вполне устраивает. Даже если для RDD алгоритм намного лучше, то разница в несколько минут и несколько процентов отклонения размеров партиций ничего не решит.
Так вот, как я вижу, будут сложности с тем, чтобы сделать класс RangePartitioner динамическим, на основании количества строк в фиксированных партициях, тем более, сначала надо объединить его с партицированием по конкретным значениям других полей.