Another fix for polars

Nixtla · Jan 10, 2025 · 555ad95 · 555ad95
1 parent 8e3e975
commit 555ad95
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/nbs/core.ipynb b/nbs/core.ipynb
@@ -78,6 +78,7 @@
     "from typing import Any, Dict, List, Optional, Sequence, Union\n",
     "\n",
     "import fsspec\n",
+    "import polars\n",
     "import numpy as np\n",
     "import pandas as pd\n",
     "import pytorch_lightning as pl\n",
@@ -1378,7 +1379,7 @@
     "\n",
     "        # Combine all series forecasts DataFrames\n",
     "        if isinstance(fcsts_dfs[0], pl_DataFrame):\n",
-    "            fcsts_df = pl.concat(fcsts_dfs, how='vertical')\n",
+    "            fcsts_df = polars.concat(fcsts_dfs, how='vertical')\n",
     "        else:\n",
     "            fcsts_df = pd.concat(fcsts_dfs, axis=0, ignore_index=True)\n",
     "        \n",
@@ -1439,7 +1440,10 @@
     "                indptr\n",
     "            )\n",
     "        # Drop duplicates when step_size < h\n",
-    "        fcsts_df = fcsts_df.drop_duplicates(subset=[self.id_col, self.time_col], keep='first')\n",
+    "        if isinstance(fcsts_df, polars.DataFrame):\n",
+    "            fcsts_df = fcsts_df.unique(subset=[self.id_col, self.time_col], keep='first')\n",
+    "        else:\n",
+    "            fcsts_df = fcsts_df.drop_duplicates(subset=[self.id_col, self.time_col], keep='first')  \n",
     "        return fcsts_df\n",
     "\n",
     "    # Save list of models with pytorch lightning save_checkpoint function\n",

diff --git a/neuralforecast/core.py b/neuralforecast/core.py
@@ -11,6 +11,7 @@
 from typing import Any, Dict, List, Optional, Sequence, Union
 
 import fsspec
+import polars
 import numpy as np
 import pandas as pd
 import pytorch_lightning as pl
@@ -1350,7 +1351,7 @@ def predict_insample(self, step_size: int = 1):
 
         # Combine all series forecasts DataFrames
         if isinstance(fcsts_dfs[0], pl_DataFrame):
-            fcsts_df = pl.concat(fcsts_dfs, how="vertical")
+            fcsts_df = polars.concat(fcsts_dfs, how="vertical")
         else:
             fcsts_df = pd.concat(fcsts_dfs, axis=0, ignore_index=True)
 
@@ -1415,9 +1416,14 @@ def predict_insample(self, step_size: int = 1):
                 fcsts_df[invert_cols].to_numpy(), indptr
             )
         # Drop duplicates when step_size < h
-        fcsts_df = fcsts_df.drop_duplicates(
-            subset=[self.id_col, self.time_col], keep="first"
-        )
+        if isinstance(fcsts_df, polars.DataFrame):
+            fcsts_df = fcsts_df.unique(
+                subset=[self.id_col, self.time_col], keep="first"
+            )
+        else:
+            fcsts_df = fcsts_df.drop_duplicates(
+                subset=[self.id_col, self.time_col], keep="first"
+            )
         return fcsts_df
 
     # Save list of models with pytorch lightning save_checkpoint function