Fixed problem with .collect and .map