persist distributed fitted values (#808)

Nixtla · Apr 4, 2024 · 8e35838 · 8e35838
1 parent 3630408
commit 8e35838
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/nbs/src/core/distributed.fugue.ipynb b/nbs/src/core/distributed.fugue.ipynb
@@ -445,7 +445,9 @@
     "                res_with_fitted = _cotransform(\n",
     "                    df, X_df, self._forecast_X_fitted, **tfm_kwargs\n",
     "                )\n",
-    "            self._results = res_with_fitted\n",
+    "            # the persist here avoids recomputing the whole thing\n",
+    "            # when retrieving the fitted values\n",
+    "            self._results = fa.persist(res_with_fitted)\n",
     "            res = transform(\n",
     "                self._results,\n",
     "                FugueBackend._retrieve_forecast_df,\n",

diff --git a/statsforecast/distributed/fugue.py b/statsforecast/distributed/fugue.py
@@ -386,7 +386,9 @@ def forecast(
                 res_with_fitted = _cotransform(
                     df, X_df, self._forecast_X_fitted, **tfm_kwargs
                 )
-            self._results = res_with_fitted
+            # the persist here avoids recomputing the whole thing
+            # when retrieving the fitted values
+            self._results = fa.persist(res_with_fitted)
             res = transform(
                 self._results,
                 FugueBackend._retrieve_forecast_df,