[FIX] Polars .columns produces list rather than Pandas Index (#862)

Nixtla · Jan 17, 2024 · e075ab8 · e075ab8
1 parent ddb03b0
commit e075ab8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/nbs/tsdataset.ipynb b/nbs/tsdataset.ipynb
@@ -354,8 +354,9 @@
     "\n",
     "        # Static features\n",
     "        if static_df is not None:\n",
-    "            static_cols = static_df.columns.drop(id_col)\n",
+    "            static_cols = [col for col in static_df.columns if col != id_col]\n",
     "            static = ufp.to_numpy(static_df[static_cols])\n",
+    "            static_cols = pd.Index(static_cols)\n",
     "        else:\n",
     "            static = None\n",
     "            static_cols = None\n",

diff --git a/neuralforecast/tsdataset.py b/neuralforecast/tsdataset.py
@@ -317,8 +317,9 @@ def from_df(
 
         # Static features
         if static_df is not None:
-            static_cols = static_df.columns.drop(id_col)
+            static_cols = [col for col in static_df.columns if col != id_col]
             static = ufp.to_numpy(static_df[static_cols])
+            static_cols = pd.Index(static_cols)
         else:
             static = None
             static_cols = None