PySpark allows Python functions to be referenced in SQL expressions via UDFs, e.g.:
df = spark.createDataFrame([Row(col_a="hello", col_b=1)])
def py_len(c):
return len(c)
spark.udf.register("pyLen", py_len)
df.selectExpr("pyLen(col_a)")
But is there a way to register a function that returns a Column expression? As expected, this doesn't work:
import pyspark.sql.functions as sqlf
def add_one(col_name: str) -> Column:
return sqlf.col(col_name) + sqlf.lit(1)
spark.udf.register("addOne", add_one)
df.selectExpr("addOne(col_b)")
Is there a way to accomplish this or something similar? This is so easy to do and such a common pattern when using the DataFrame API that I thought it would be simple in pure SQL, but it doesn't seem that way.