nixtla_fugue_spark.py

from statsforecast.distributed.utils import forecast
from statsforecast.distributed.fugue import FugueBackend
from statsforecast.models import AutoARIMA
from statsforecast.core import StatsForecast

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
backend = FugueBackend(spark, {"fugue.spark.use_pandas_udf":True})

forecast(spark.read.parquet("/tmp/1m.parquet"), 
         [AutoARIMA()], 
         freq="D", 
         h=7, 
         parallel=backend).toPandas()