[SPARK-3141] [PySpark] fix sortByKey() with take()

Fix sortByKey() with take()

The function `f` used in mapPartitions should always return an iterator.