Ошибка Py4j при использовании функции Spark UDF

Question

Я пытаюсь создать UDF для использования его в запросе spark sql с именами столбцов. Я получаю Py4JJavaError.

Пожалуйста, помогите мне решить проблему. Я хочу использовать dataframe.show(), чтобы увидеть результат.

display(ss) and ss.printSchema() работают, но ss.show() не работает.

from pyspark.sql.types import IntegerType

def squared_typed(s):
    if not s: return
    return int(s * s)
spark.udf.register("squaredWithPython", squared_typed, IntegerType())

spark.range(1, 20).registerTempTable("test")
ss = spark.sql('''select squaredWithPython(id) as idsquared from test''')