from pyspark.sql.types import LongType def squared_typed(s): return s * s spark.udf.register("squaredWithPython", squared_typed, LongType())