nixtla_fugue.py

from fugue import transform

def forecast_series(df: pd.DataFrame, models) -> pd.DataFrame:
    tdf = df.set_index("unique_id")
    model = StatsForecast(df=tdf, models=models, freq='D', n_jobs=1)
    return model.forecast(7).reset_index()

transform(series.reset_index(),
          forecast_series,
          params=dict(models=[AutoARIMA()]),
          schema="unique_id:int, ds:date, AutoARIMA:float",
          partition={"by": "unique_id"},
          engine="spark"
          ).show()