Structured streaming multiple row to pandas udf - pandas

Question

I'm writing a structured streaming job that receives data from eventhubs.
After some preparation, I apply a pandas_udf function on each row to create a new column with a prediction from a pickle model.
I'm experiencing a serious problem: sometimes the input for the pandas_udf is a group of row and not a single row (as expected). This leads me to an error:
RuntimeError: Result vector from pandas_udf was not the required length: expected 2, got 1
This happens because the pandas_udf receives more than one row (in this case 2).
How could this be possible? Shouldn't the .withColumn be executed row-wise on each row?
Here is my code:
dfInt = spark \
.readStream \
.load() \
.selectExpr("cast (body as string) as json") \
.select(from_json("json",schema).alias("data")) \
.withColumn("k", expr("uuid()")) \
.select("key", explode("data.features").alias("feat")) \
.select("feat.*", "key") \
.groupBy("k") \
.agg(*expressions) \
.drop("k") \
.na.drop() \
.withColumn("prediction", predict( (F.struct([col(x) for x in (features)]))))
The pandas_udf is the following:
#pandas_udf(FloatType())
def predict(x):
return pd.Series(pickle_model.predict_proba(x)[0][1])
Actually the problems seems to be before the withColumn call with the udf, because more row are coming from the previous step.
The groupBy aggregation returns a singular row, because the key on which I make the group by is unique.
Do you know which is the reason for that?

Answer 1

In this case you are using a SCALAR pandas_udf, which takes as input a pandas Series and returns a pandas.Series of the same size. I don't know the exact details on the internals but my understanding is that each executor will convert your column (F.struct([col(x) for x in (features)])) into a pandas.Series for the Dataframe partition the executor is currently processing and apply the function on the series. A partition consists of many rows, therefore you cannot assume that the series is only of length one. You need to make sure that you are keeping all your predicted proba for all your rows. You can probably do something like this (assuming you are indeed only interested in keeping the probability of class 1):
#pandas_udf(FloatType())
def predict(x):
return pd.Series(pickle_model.predict_proba(x)[:,1])