我有一个spark.DataFrame，有多个时间序列。我想在groupby apply上为每个时间序列应用一个sklearn模型。对于每个时间序列，该模型需要大约0.05s，但是当我尝试在pandas_udf中求解这个问题时，它比顺序应用要花费更长的时间。下面是一个示例

def forecaster_spark(data_group: pd.DataFrame):
    # index for reports
    item_id = data_group["item_id"].iloc[0]
    # Indexing by time index
    data_group = data_group.set_index(pd.DatetimeIndex(data_group['ds'])).sort_index()
    # Here we extract the time series
    y = data_group["y"].astype(float)
    # transform is a transformation to build regressor features (example extracting lags)
    X = transform(y)
    model = xg.XGBRegressor(max_depth = 50)
    # For each item, the model takes app 0.05s
    model.fit(X[:-1],y[-2])
    return pd.DataFrame({"item_id": item_id, "pred": model.predict(X.iloc[-1], y[-1])})

这里，transform是应用于y的转换，用于构建回归矩阵，例如，可以将其视为构建滞后的转换。
然后，我在包含所有数据的spark.DataFrame上应用此函数

predictions = data.applyInPandas(forecaster_spark, schema="item_id string, y_pred double")

在predictions上执行操作时（例如toPandas或show(1000)），这会花费很多时间。当我在forecaster_spark上注解model.fit时（改变pred返回任何值）它完成得非常快，所以问题是在UDF内部拟合模型。我想可能是使用groupby时 Dataframe 的分区有问题，但是我尝试了许多简单的UDF函数（例如取平均值），结果都很正常。

def test_function_for_udf(data_group: pd.DataFrame):
    """Same function as above, but commenting fit method"""
    # index for reports
    item_id = data_group["item_id"].iloc[0]
    # Indexing by time index
    data_group = data_group.set_index(pd.DatetimeIndex(data_group['ds'])).sort_index()
    # Here we extract the time series
    y = data_group["y"].astype(float)
    # transform is a transformation to build regressor features (example extracting lags)
    X = transform(y)
    model = xg.XGBRegressor(max_depth = 50)
    # Note that I commented the fit
    # model.fit(X[:-1],y[-2])
    return pd.DataFrame({"item_id": item_id, "mean": y.mean()})