When to use `pandas_udf` of iterator of Series to iterator of Series in pyspark?

Viewed 395

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

  1. 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|
+-----------+
  1. 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:

  1. In plus_one(), it prints y as much as the length of the batch_iter, which is a sort of contradiction to useful when the UDF execution requires initializing some state. I don't see why it gives benefit in this case.
  2. pandas_udf of Series -> Series and pandas_udf of Iterator[Series] -> Iterator[Series] work SAME. I think there is no difference between the way they work?
0 Answers
Related