Не работает DecisionTreeClassifier (pyspark.ml, Spark 2.4.0)

В моей таблице около 50 млн. записей, около 400 столбцов. У меня не получается использовать DecisionTreeClassifier (pyspark.ml, Spark 2.4.0) - всё зависает, хотя для маленьких выборок всё работает. Модель я делал примерно как здесь: https://habr.com/ru/post/522296/. Подскажите, пожалуйста, что можно сделать.