fix(distributed): exogenous handling in distributed cross validation (#443)

Author: jmoralez

mlforecast/distributed/forecast.py

@@ -195,7 +195,7 @@ def _preprocess_partition(
             train = part[train_mask]
             valid_keep_cols = part.columns
             if static_features is not None:
-                valid_keep_cols.drop(static_features)
+                valid_keep_cols = valid_keep_cols.drop(static_features)
             valid = part.loc[valid_mask, valid_keep_cols].merge(cutoffs, on=id_col)
         transformed = ts.fit_transform(
             train,
@@ -456,6 +456,8 @@ def _predict(
     ) -> Iterable[pd.DataFrame]:
         for serialized_ts, _, serialized_valid in items:
             valid = cloudpickle.loads(serialized_valid)
+            if valid is not None:
+                X_df = valid
             ts = cloudpickle.loads(serialized_ts)
             res = ts.predict(
                 models=models,
@@ -649,7 +651,11 @@ def cross_validation(
                 engine=self.engine,
             )
             results.append(fa.get_native_as_df(preds))
-        return fa.union(*results)
+        if len(results) == 1:
+            res = results[0]
+        else:
+            res = fa.union(*results)
+        return res
 
     @staticmethod
     def _save_ts(items: List[List[Any]], path: str) -> Iterable[pd.DataFrame]:

nbs/distributed.forecast.ipynb

@@ -261,7 +261,7 @@
     "            train = part[train_mask]\n",
     "            valid_keep_cols = part.columns\n",
     "            if static_features is not None:\n",
-    "                valid_keep_cols.drop(static_features)\n",
+    "                valid_keep_cols = valid_keep_cols.drop(static_features)\n",
     "            valid = part.loc[valid_mask, valid_keep_cols].merge(cutoffs, on=id_col)\n",
     "        transformed = ts.fit_transform(\n",
     "            train,\n",
@@ -508,6 +508,8 @@
     "    ) -> Iterable[pd.DataFrame]:\n",
     "        for serialized_ts, _, serialized_valid in items:\n",
     "            valid = cloudpickle.loads(serialized_valid)\n",
+    "            if valid is not None:\n",
+    "                X_df = valid\n",
     "            ts = cloudpickle.loads(serialized_ts)\n",
     "            res = ts.predict(\n",
     "                models=models,\n",
@@ -695,7 +697,11 @@
     "                engine=self.engine,\n",
     "            )\n",
     "            results.append(fa.get_native_as_df(preds))\n",
-    "        return fa.union(*results)\n",
+    "        if len(results) == 1:\n",
+    "            res = results[0]\n",
+    "        else:\n",
+    "            res = fa.union(*results)\n",
+    "        return res\n",
     "\n",
     "    @staticmethod\n",
     "    def _save_ts(items: List[List[Any]], path: str) -> Iterable[pd.DataFrame]:\n",