Enable mypy check back with PySpark 3.1. #2088

PySpark 3.1 ported pyspark-stubs into its codebase which causes mypy check errors.
We disabled it at #2083 but we should enable it back as soon as possible.