Fixes for low versions of pandas

tinkoff-ai · Mar 29, 2023 · 8e169ef · 8e169ef
1 parent 5415a78
commit 8e169ef
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/etna/datasets/tsdataset.py b/etna/datasets/tsdataset.py
@@ -1135,7 +1135,8 @@ def get_level_dataset(self, target_level: str) -> "TSDataset":
             target_level_df = self.to_pandas(features=target_names)
 
         target_components_df = target_level_df.loc[:, pd.IndexSlice[:, self.target_components_names]]
-        target_level_df = target_level_df.drop(columns=list(self.target_components_names), level="feature")
+        if len(self.target_components_names) > 0:  # for pandas >=1.1, <1.2
+            target_level_df = target_level_df.drop(columns=list(self.target_components_names), level="feature")
 
         ts = TSDataset(
             df=target_level_df,
@@ -1202,8 +1203,9 @@ def get_target_components(self) -> Optional[pd.DataFrame]:
 
     def drop_target_components(self):
         """Drop target components from dataset."""
-        self.df.drop(columns=list(self.target_components_names), level="feature", inplace=True)
-        self._target_components_names = ()
+        if len(self.target_components_names) > 0:  # for pandas >=1.1, <1.2
+            self.df.drop(columns=list(self.target_components_names), level="feature", inplace=True)
+            self._target_components_names = ()
 
     @property
     def columns(self) -> pd.core.indexes.multi.MultiIndex: