Fix Catboost prediction_components

etna/models/catboost.py

@@ -145,8 +145,10 @@ def predict_components(self, df: pd.DataFrame) -> pd.DataFrame:
             dataframe with prediction components
         """
         features = df.drop(columns=["timestamp", "target"])
+        self._prepare_float_category_columns(features)
+        predict_pool = Pool(features, cat_features=self._categorical)
+        prediction = self.model.predict(predict_pool)
 
-        prediction = self.model.predict(features)
         pool = self._prepare_pool(features, prediction)
         shap_values = self.model.get_feature_importance(pool, type="ShapValues")
 

tests/test_models/test_catboost.py

@@ -183,9 +183,23 @@ def test_decomposition_sums_to_target(dfs_w_exog):
     np.testing.assert_allclose(y_hat_pred, y_pred)
 
 
+@pytest.fixture
+def ts_with_features() -> TSDataset:
+    df = generate_ar_df(periods=100, n_segments=2, start_time="2000-01-01")
+    df_exog = df.copy().rename(columns={"target": "exog"})
+    df_exog["cat_exog"] = df_exog["exog"].astype(int).astype("category")
+
+    df = TSDataset.to_dataset(df)
+    df_exog = TSDataset.to_dataset(df_exog)
+    return TSDataset(df=df, df_exog=df_exog, freq="D")
+
+
 @pytest.mark.parametrize("model", (CatBoostPerSegmentModel(), CatBoostMultiSegmentModel()))
-def test_prediction_decomposition(outliers_tsds, model):
-    train, test = outliers_tsds.train_test_split(test_size=10)
+def test_prediction_decomposition(ts_with_features, model):
+    train, test = ts_with_features.train_test_split(test_size=10)
+    test.df = test.df.loc[
+        pd.IndexSlice[:], pd.IndexSlice[:, ["target", "exog", "cat_exog"]]
+    ]  # Order differs from train
     assert_prediction_components_are_present(model=model, train=train, test=test)