Pyspark 2.4.1 с pandas udf показывает странный вывод

Question

Я проверяю функциональность pandas_udf для сгруппированной карты с pyspark, которая запускается как PYSPARK_PYTHON=python3 pyspark.

Python version: 3.6.8
pyarrow: 0.13.0
pyspark: 2.4.1

Это код:

import numpy as np
import pandas as pd
spark.conf.set("spark.sql.execution.arrow.enabled", "true")

from pyspark.sql.functions import pandas_udf, PandasUDFType

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))

@pandas_udf("id long, v double", PandasUDFType.GROUPED_MAP)
def subtract_mean(pdf):
    # pdf is a pandas.DataFrame
    v = pdf.v
    return pdf.assign(v=v - v.mean())

df.groupby("id").apply(subtract_mean).show()

Результатожидается как указано ниже в соответствии с Руководством по использованию PySpark для панд с Apache Arrow .

# +---+----+
# | id|   v|
# +---+----+
# |  1|-0.5|
# |  1| 0.5|
# |  2|-3.0|
# |  2|-1.0|
# |  2| 4.0|
# +---+----+

Однако у меня есть

+----+----+
|  id|   v|
+----+----+
|   1| 0.5|
|null|null|
|   2|-1.0|
|   2| 4.0|
|null|null|
+----+----+

Однако этот код работает какожидается в Spark 2.3.1.Я что-то не так или это ошибка?