Repartition to 1 mln without logic it has to be slow... and then collect() which run on driver and shouldn't be run in production.

That code divides your dataset into one mln parts on workers (also causes unnecessary data exchange on workers), then transfers it back to the driver and merges them for the collect function. Esing RDD is also superfluous (you don't benefit from Adaptive Query Execution when you use RDD directly)