Fix TSDataset.make_future to handle hierarchy, quantiles, target components

CHANGELOG.md

@@ -19,7 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 -
 - Fix `BaseReconciliator` to work on `pandas==1.1.5` ([#1229](https://github.com/tinkoff-ai/etna/pull/1229))
--
+- Fix `TSDataset.make_future` to handle hierarchy, quantiles, target components ([#1248](https://github.com/tinkoff-ai/etna/pull/1248))
 - Fix warning during creation of `ResampleWithDistributionTransform` ([#1230](https://github.com/tinkoff-ai/etna/pull/1230))
 - Add deep copy for copying attributes of `TSDataset` ([#1241](https://github.com/tinkoff-ai/etna/pull/1241))
 -

etna/datasets/tsdataset.py

@@ -230,21 +230,23 @@ def __getitem__(self, item):
     def make_future(
         self, future_steps: int, transforms: Sequence["Transform"] = (), tail_steps: int = 0
     ) -> "TSDataset":
-        """Return new TSDataset with future steps.
+        """Return new TSDataset with features extended into the future.
+
+        The result dataset doesn't contain quantiles and target components.
 
         Parameters
         ----------
         future_steps:
-            number of timestamp in the future to build features for.
+            number of steps to extend dataset into the future.
         transforms:
             sequence of transforms to be applied.
         tail_steps:
-            number of timestamp for context to build features for.
+            number of steps to keep from the tail of the original dataset.
 
         Returns
         -------
         :
-            dataset with features in the future.
+            dataset with features extended into the.
 
         Examples
         --------
@@ -295,6 +297,14 @@ def make_future(
                             f"NaN-s will be used for missing values"
                         )
 
+        # remove components and quantiles
+        # it should be done if we have quantiles and components in raw_df
+        # TODO: fix this after making quantiles to work like components, with special methods
+        if len(self.target_components_names) > 0:
+            df = df.drop(columns=list(self.target_components_names), level="feature")
+        if len(self.target_quantiles_names) > 0:
+            df = df.drop(columns=list(self.target_quantiles_names), level="feature")
+
         # Here only df is required, other metadata is not necessary to build the dataset
         ts = TSDataset(df=df, freq=self.freq)
         for transform in transforms:
@@ -305,7 +315,7 @@ def make_future(
         future_dataset = df.tail(future_steps + tail_steps).copy(deep=True)
 
         future_dataset = future_dataset.sort_index(axis=1, level=(0, 1))
-        future_ts = TSDataset(df=future_dataset, freq=self.freq)
+        future_ts = TSDataset(df=future_dataset, freq=self.freq, hierarchical_structure=self.hierarchical_structure)
 
         # can't put known_future into constructor, _check_known_future fails with df_exog=None
         future_ts.known_future = deepcopy(self.known_future)

tests/test_datasets/test_dataset.py

@@ -620,6 +620,24 @@ def test_make_future_inherits_regressors(df_and_regressors):
     assert ts_future.regressors == ts.regressors
 
 
+def test_make_future_inherits_hierarchy(product_level_constant_forecast_with_quantiles):
+    ts = product_level_constant_forecast_with_quantiles
+    future = ts.make_future(future_steps=2)
+    assert future.hierarchical_structure is ts.hierarchical_structure
+
+
+def test_make_future_removes_quantiles(product_level_constant_forecast_with_quantiles):
+    ts = product_level_constant_forecast_with_quantiles
+    future = ts.make_future(future_steps=2)
+    assert len(future.target_quantiles_names) == 0
+
+
+def test_make_future_removes_target_components(ts_with_target_components):
+    ts = ts_with_target_components
+    future = ts.make_future(future_steps=2)
+    assert len(future.target_components_names) == 0
+
+
 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