In this page: This pandas UDF is useful when the UDF execution requires initializing some state, for example, loading a machine learning model file to apply inference to every input batch.
So, I've tested it:
import pandas as pd
from typing import Iterator
from pyspark.sql.functions import col, pandas_udf
pdf = pd.DataFrame(range(50), columns=["x"])
df = spark.createDataFrame(pdf)
y_bc = spark.sparkContext.broadcast(99999)
#
# Method1: Iterator[pd.Series] -> Iterator[pd.Series]
#
@pandas_udf("long")
def plus_one(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
y = y_bc.value
print(y) # Check how many times it print it out
for x in batch_iter:
print(x)
yield x + 1
#
# Method2: pd.Series -> pd.Series
#
@pandas_udf("long")
def plus_one2(v: pd.Series) -> pd.Series:
y = y_bc.value
print(y) # Check how many times it print it out
return v + 1
df.select(plus_one(col("x"))).show()
df.select(plus_one2(col("x"))).show()
Result
df.select(plus_one(col("x"))).show()
99999
99999
99999
99999
99999
99999
99999
99999
+-----------+
|plus_one(x)|
+-----------+
| 1|
| 2|
| 3|
| 4|
| 5|
| 6|
| 7|
| 8|
| 9|
| 10|
| 11|
| 12|
| 13|
| 14|
| 15|
| 16|
| 17|
| 18|
| 19|
| 20|
+-----------+
df.select(plus_one2(col("x"))).show()
99999
99999
99999
99999
99999
99999
99999
99999
+------------+
|plus_one2(x)|
+------------+
| 1|
| 2|
| 3|
| 4|
| 5|
| 6|
| 7|
| 8|
| 9|
| 10|
| 11|
| 12|
| 13|
| 14|
| 15|
| 16|
| 17|
| 18|
| 19|
| 20|
+------------+
only showing top 20 rows
Result analysis:
- In
plus_one(), it printsyas much as the length of thebatch_iter, which is a sort of contradiction touseful when the UDF execution requires initializing some state. I don't see why it gives benefit in this case. pandas_udfofSeries -> Seriesandpandas_udfofIterator[Series] -> Iterator[Series]work SAME. I think there is no difference between the way they work?