Update make_future and train_test_split regressors handling

etna/datasets/tsdataset.py

@@ -190,6 +190,7 @@ def _update_regressors(self, transform: "Transform", columns_before: Set[str], c
         else:
             raise ValueError("Transform is not FutureMixin and does not have in_column attribute!")
 
+        new_regressors = [regressor for regressor in new_regressors if regressor not in self.regressors]
         self._regressors.extend(new_regressors)
 
     def __repr__(self):
@@ -278,7 +279,9 @@ def make_future(self, future_steps: int) -> "TSDataset":
 
         future_dataset = df.tail(future_steps).copy(deep=True)
         future_dataset = future_dataset.sort_index(axis=1, level=(0, 1))
-        future_ts = TSDataset(future_dataset, freq=self.freq)
+        future_ts = TSDataset(df=future_dataset, freq=self.freq)
+        future_ts.known_future = self.known_future
+        future_ts._regressors = self.regressors
         future_ts.transforms = self.transforms
         future_ts.df_exog = self.df_exog
         return future_ts
@@ -332,7 +335,7 @@ def _check_regressors(df: pd.DataFrame, df_regressors: pd.DataFrame):
 
     def _merge_exog(self, df: pd.DataFrame) -> pd.DataFrame:
         segments = sorted(set(df.columns.get_level_values("segment")))
-        df_regressors = self.df_exog.loc[:, pd.IndexSlice[segments, self.known_future]]
+        df_regressors = self.df_exog.loc[:, pd.IndexSlice[segments, self.regressors]]
         self._check_regressors(df=df, df_regressors=df_regressors)
         df = pd.merge(df, self.df_exog, left_index=True, right_index=True, how="left").sort_index(axis=1, level=(0, 1))
         return df
@@ -730,13 +733,15 @@ def train_test_split(
 
         train_df = self.df[train_start_defined:train_end_defined][self.raw_df.columns]  # type: ignore
         train_raw_df = self.raw_df[train_start_defined:train_end_defined]  # type: ignore
-        train = TSDataset(df=train_df, df_exog=self.df_exog, freq=self.freq)
+        train = TSDataset(df=train_df, df_exog=self.df_exog, freq=self.freq, known_future=self.known_future)
         train.raw_df = train_raw_df
+        train._regressors = self.regressors
 
         test_df = self.df[test_start_defined:test_end_defined][self.raw_df.columns]  # type: ignore
         test_raw_df = self.raw_df[train_start_defined:test_end_defined]  # type: ignore
-        test = TSDataset(df=test_df, df_exog=self.df_exog, freq=self.freq)
+        test = TSDataset(df=test_df, df_exog=self.df_exog, freq=self.freq, known_future=self.known_future)
         test.raw_df = test_raw_df
+        test._regressors = self.regressors
 
         return train, test
 

etna/pipeline/pipeline.py

@@ -112,7 +112,7 @@ def fit(self, ts: TSDataset) -> "Pipeline":
 
     def _forecast_prediction_interval(self, future: TSDataset) -> TSDataset:
         """Forecast prediction interval for the future."""
-        _, forecasts, _ = self.backtest(self.ts, metrics=[MAE()], n_folds=self.n_folds)
+        _, forecasts, _ = self.backtest(ts=self.ts, metrics=[MAE()], n_folds=self.n_folds)
         forecasts = TSDataset(df=forecasts, freq=self.ts.freq)
         residuals = (
             forecasts.loc[:, pd.IndexSlice[:, "target"]]

tests/test_datasets/test_dataset.py

@@ -289,6 +289,14 @@ def test_train_test_split_failed(test_size, borders, match, tsdf_with_exog):
         )
 
 
+def test_train_test_split_pass_regressors_to_output(df_and_regressors):
+    df, df_exog, known_future = df_and_regressors
+    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
+    train, test = ts.train_test_split(test_size=5)
+    assert train.regressors == ts.regressors
+    assert test.regressors == ts.regressors
+
+
 def test_dataset_datetime_conversion():
     classic_df = generate_ar_df(periods=30, start_time="2021-06-01", n_segments=2)
     classic_df["timestamp"] = classic_df["timestamp"].astype(str)
@@ -355,6 +363,13 @@ def test_make_future_with_regressors(df_and_regressors):
     assert set(ts_future.columns.get_level_values("feature")) == {"target", "regressor_1", "regressor_2"}
 
 
+def test_make_future_inherits_regressors(df_and_regressors):
+    df, df_exog, known_future = df_and_regressors
+    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
+    ts_future = ts.make_future(10)
+    assert ts_future.regressors == ts.regressors
+
+
 def test_make_future_warn_not_enough_regressors(df_and_regressors):
     """Check that warning is thrown if regressors don't have enough values for the future."""
     df, df_exog, known_future = df_and_regressors