From 71b80d2d7e1ec0f44391f9a7768435da6f0ab9cd Mon Sep 17 00:00:00 2001
From: Alexander76Kuznetsov
 <128227577+Alexander76Kuznetsov@users.noreply.github.com>
Date: Mon, 19 Feb 2024 17:44:02 +0300
Subject: [PATCH] `StandardScalerTransform` optimization (#236)

---
 etna/transforms/math/sklearn.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/etna/transforms/math/sklearn.py b/etna/transforms/math/sklearn.py
index d038ebdf4..f09af45aa 100644
--- a/etna/transforms/math/sklearn.py
+++ b/etna/transforms/math/sklearn.py
@@ -3,6 +3,7 @@
 from typing import Dict
 from typing import List
 from typing import Optional
+from typing import Sequence
 from typing import Union
 from typing import cast
 
@@ -226,9 +227,8 @@ def _inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame:
         return df
 
     def _preprocess_macro(self, df: pd.DataFrame) -> np.ndarray:
-        segments = sorted(set(df.columns.get_level_values("segment")))
-        x = df.loc[:, pd.IndexSlice[:, self.in_column]]
-        x = pd.concat([x[segment] for segment in segments]).values
+        in_column = cast(Sequence[str], self.in_column)
+        x = TSDataset.to_flatten(df, features=in_column)[self.in_column].values
         return x
 
     def _postprocess_macro(self, df: pd.DataFrame, transformed: np.ndarray) -> np.ndarray: