Вместо датасетов с большим числом колонок лучше делать колонку-вектор или колонку-мапу. Собственно если это результат пивота то сам бог велел OneHotEncoder сделать. Если при обучении регрессий провал то, скорее всего там мелкие задачи и идет много накладных расходов на драйвер - надо попробовать число партиций уменьшить и подумать что и в какой момент подкешить.
Сперва нужно отпроцессить длинный фрейм, а потом спивотить (и там происходила посадка из-за большого числа столбцов в результате пивота, но удалось ускорить) тут, к сожалению, без вариантов. Про ванхот не понял, это не категории, зачем он? Да, по поводу мелких задач согласен, но по партициям там и так их всего 5. Я выше код на пастебин прикладывал, там все кэшируется, я это усвоил когда в логах спарк варнинги кидал про перфоманс декриз если не будет кэша))