upgini 1.1.280.dev0__py3-none-any.whl → 1.2.31a2__py3-none-any.whl

upgini/features_enricher.py

@@ -11,6 +11,7 @@ import sys
 import tempfile
 import time
 import uuid
+from collections import Counter
 from dataclasses import dataclass
 from threading import Thread
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
@@ -22,7 +23,6 @@ from pandas.api.types import (
     is_datetime64_any_dtype,
     is_numeric_dtype,
     is_object_dtype,
-    is_period_dtype,
     is_string_dtype,
 )
 from scipy.stats import ks_2samp
@@ -45,24 +45,31 @@ from upgini.mdc import MDC
 from upgini.metadata import (
     COUNTRY,
     DEFAULT_INDEX,
+    ENTITY_SYSTEM_RECORD_ID,
     EVAL_SET_INDEX,
     ORIGINAL_INDEX,
     RENAMED_INDEX,
+    SEARCH_KEY_UNNEST,
     SORT_ID,
     SYSTEM_RECORD_ID,
     TARGET,
     CVType,
+    FeaturesMetadataV2,
     FileColumnMeaningType,
     ModelTaskType,
     RuntimeParameters,
     SearchKey,
 )
 from upgini.metrics import EstimatorWrapper, validate_scoring_argument
+from upgini.normalizer.normalize_utils import Normalizer
 from upgini.resource_bundle import ResourceBundle, bundle, get_custom_bundle
 from upgini.search_task import SearchTask
 from upgini.spinner import Spinner
 from upgini.utils import combine_search_keys, find_numbers_with_decimal_comma
-from upgini.utils.country_utils import CountrySearchKeyDetector
+from upgini.utils.country_utils import (
+    CountrySearchKeyConverter,
+    CountrySearchKeyDetector,
+)
 from upgini.utils.custom_loss_utils import (
     get_additional_params_custom_loss,
     get_runtime_params_custom_loss,
@@ -71,8 +78,8 @@ from upgini.utils.cv_utils import CVConfig, get_groups
 from upgini.utils.datetime_utils import (
     DateTimeSearchKeyConverter,
     is_blocked_time_series,
+    is_dates_distribution_valid,
     is_time_series,
-    validate_dates_distribution,
 )
 from upgini.utils.deduplicate_utils import (
     clean_full_duplicates,
@@ -84,12 +91,20 @@ from upgini.utils.display_utils import (
     prepare_and_show_report,
     show_request_quote_button,
 )
-from upgini.utils.email_utils import EmailSearchKeyConverter, EmailSearchKeyDetector
+from upgini.utils.email_utils import (
+    EmailDomainGenerator,
+    EmailSearchKeyConverter,
+    EmailSearchKeyDetector,
+)
+from upgini.utils.feature_info import FeatureInfo, _round_shap_value
 from upgini.utils.features_validator import FeaturesValidator
 from upgini.utils.format import Format
-from upgini.utils.ip_utils import IpToCountrySearchKeyConverter
-from upgini.utils.phone_utils import PhoneSearchKeyDetector
-from upgini.utils.postal_code_utils import PostalCodeSearchKeyDetector
+from upgini.utils.ip_utils import IpSearchKeyConverter
+from upgini.utils.phone_utils import PhoneSearchKeyConverter, PhoneSearchKeyDetector
+from upgini.utils.postal_code_utils import (
+    PostalCodeSearchKeyConverter,
+    PostalCodeSearchKeyDetector,
+)
 
 try:
     from upgini.utils.progress_bar import CustomProgressBar as ProgressBar
@@ -145,6 +160,10 @@ class FeaturesEnricher(TransformerMixin):
 
     shared_datasets: list of str, optional (default=None)
         List of private shared dataset ids for custom search
+
+    select_features: bool, optional (default=False)
+        If True, return only selected features both from input and data sources.
+        Otherwise, return all features from input and only selected features from data sources.
     """
 
     TARGET_NAME = "target"
@@ -211,11 +230,12 @@ class FeaturesEnricher(TransformerMixin):
         client_visitorid: Optional[str] = None,
         custom_bundle_config: Optional[str] = None,
         add_date_if_missing: bool = True,
+        select_features: bool = False,
         **kwargs,
     ):
         self.bundle = get_custom_bundle(custom_bundle_config)
         self._api_key = api_key or os.environ.get(UPGINI_API_KEY)
-        if api_key is not None and not isinstance(api_key, str):
+        if self._api_key is not None and not isinstance(self._api_key, str):
             raise ValidationError(f"api_key should be `string`, but passed: `{api_key}`")
         self.rest_client = get_rest_client(endpoint, self._api_key, client_ip, client_visitorid)
         self.client_ip = client_ip
@@ -235,6 +255,7 @@ class FeaturesEnricher(TransformerMixin):
 
         self.passed_features: List[str] = []
         self.df_with_original_index: Optional[pd.DataFrame] = None
+        self.fit_columns_renaming: Optional[Dict[str, str]] = None
         self.country_added = False
         self.fit_generated_features: List[str] = []
         self.fit_dropped_features: Set[str] = set()
@@ -245,10 +266,12 @@ class FeaturesEnricher(TransformerMixin):
         self.eval_set: Optional[List[Tuple]] = None
         self.autodetected_search_keys: Dict[str, SearchKey] = {}
         self.imbalanced = False
-        self.__cached_sampled_datasets: Optional[Tuple[pd.DataFrame, pd.DataFrame, pd.Series, Dict, Dict]] = None
+        self.__cached_sampled_datasets: Dict[str, Tuple[pd.DataFrame, pd.DataFrame, pd.Series, Dict, Dict, Dict]] = (
+            dict()
+        )
 
-        validate_version(self.logger)
-        self.search_keys = search_keys or dict()
+        validate_version(self.logger, self.__log_warning)
+        self.search_keys = search_keys or {}
         self.country_code = country_code
         self.__validate_search_keys(search_keys, search_id)
         self.model_task_type = model_task_type
@@ -261,8 +284,11 @@ class FeaturesEnricher(TransformerMixin):
         self._relevant_data_sources_wo_links: pd.DataFrame = self.EMPTY_DATA_SOURCES
         self.metrics: Optional[pd.DataFrame] = None
         self.feature_names_ = []
+        self.dropped_client_feature_names_ = []
         self.feature_importances_ = []
         self.search_id = search_id
+        self.select_features = select_features
+
         if search_id:
             search_task = SearchTask(search_id, rest_client=self.rest_client, logger=self.logger)
 
@@ -322,6 +348,10 @@ class FeaturesEnricher(TransformerMixin):
         self.exclude_columns = exclude_columns
         self.baseline_score_column = baseline_score_column
         self.add_date_if_missing = add_date_if_missing
+        self.features_info_display_handle = None
+        self.data_sources_display_handle = None
+        self.autofe_features_display_handle = None
+        self.report_button_handle = None
 
     def _get_api_key(self):
         return self._api_key
@@ -423,7 +453,7 @@ class FeaturesEnricher(TransformerMixin):
 
             self.logger.info("Start fit")
 
-            self.__validate_search_keys(self.search_keys, self.search_id)
+            self.__validate_search_keys(self.search_keys)
 
             # Validate client estimator params
             self._get_client_cat_features(estimator, X, self.search_keys)
@@ -557,7 +587,7 @@ class FeaturesEnricher(TransformerMixin):
 
             self.logger.info("Start fit_transform")
 
-            self.__validate_search_keys(self.search_keys, self.search_id)
+            self.__validate_search_keys(self.search_keys)
 
             search_progress = SearchProgress(0.0, ProgressStage.START_FIT)
             if progress_callback is not None:
@@ -704,7 +734,7 @@ class FeaturesEnricher(TransformerMixin):
 
             start_time = time.time()
             try:
-                result = self.__inner_transform(
+                result, _, _ = self.__inner_transform(
                     trace_id,
                     X,
                     exclude_features_sources=exclude_features_sources,
@@ -831,17 +861,44 @@ class FeaturesEnricher(TransformerMixin):
                 self.logger.warning(msg)
                 print(msg)
 
+            if X is not None and y is None:
+                raise ValidationError("X passed without y")
+
             self.__validate_search_keys(self.search_keys, self.search_id)
             effective_X = X if X is not None else self.X
             effective_y = y if y is not None else self.y
             effective_eval_set = eval_set if eval_set is not None else self.eval_set
             effective_eval_set = self._check_eval_set(effective_eval_set, effective_X, self.bundle)
 
+            if (
+                self._search_task is None
+                or self._search_task.provider_metadata_v2 is None
+                or len(self._search_task.provider_metadata_v2) == 0
+                or effective_X is None
+                or effective_y is None
+            ):
+                raise ValidationError(self.bundle.get("metrics_unfitted_enricher"))
+
+            validated_X = self._validate_X(effective_X)
+            validated_y = self._validate_y(validated_X, effective_y)
+            validated_eval_set = (
+                [self._validate_eval_set_pair(validated_X, eval_pair) for eval_pair in effective_eval_set]
+                if effective_eval_set is not None
+                else None
+            )
+
+            if self.X is None:
+                self.X = X
+            if self.y is None:
+                self.y = y
+            if self.eval_set is None:
+                self.eval_set = effective_eval_set
+
             try:
                 self.__log_debug_information(
-                    effective_X,
-                    effective_y,
-                    effective_eval_set,
+                    validated_X,
+                    validated_y,
+                    validated_eval_set,
                     exclude_features_sources=exclude_features_sources,
                     cv=cv if cv is not None else self.cv,
                     importance_threshold=importance_threshold,
@@ -851,21 +908,9 @@ class FeaturesEnricher(TransformerMixin):
                     remove_outliers_calc_metrics=remove_outliers_calc_metrics,
                 )
 
-                if (
-                    self._search_task is None
-                    or self._search_task.provider_metadata_v2 is None
-                    or len(self._search_task.provider_metadata_v2) == 0
-                    or effective_X is None
-                    or effective_y is None
-                ):
-                    raise ValidationError(self.bundle.get("metrics_unfitted_enricher"))
-
-                if X is not None and y is None:
-                    raise ValidationError("X passed without y")
-
                 validate_scoring_argument(scoring)
 
-                self._validate_baseline_score(effective_X, effective_eval_set)
+                self._validate_baseline_score(validated_X, validated_eval_set)
 
                 if self._has_paid_features(exclude_features_sources):
                     msg = self.bundle.get("metrics_with_paid_features")
@@ -874,14 +919,14 @@ class FeaturesEnricher(TransformerMixin):
                     return None
 
                 cat_features, search_keys_for_metrics = self._get_client_cat_features(
-                    estimator, effective_X, self.search_keys
+                    estimator, validated_X, self.search_keys
                 )
 
                 prepared_data = self._prepare_data_for_metrics(
                     trace_id=trace_id,
-                    X=effective_X,
-                    y=effective_y,
-                    eval_set=effective_eval_set,
+                    X=X,
+                    y=y,
+                    eval_set=eval_set,
                     exclude_features_sources=exclude_features_sources,
                     importance_threshold=importance_threshold,
                     max_features=max_features,
@@ -904,21 +949,27 @@ class FeaturesEnricher(TransformerMixin):
                     search_keys,
                     groups,
                     _cv,
+                    columns_renaming,
                 ) = prepared_data
 
+                # rename cat_features
+                if cat_features:
+                    for new_c, old_c in columns_renaming.items():
+                        if old_c in cat_features:
+                            cat_features.remove(old_c)
+                            cat_features.append(new_c)
+
                 gc.collect()
 
+                if fitting_X.shape[1] == 0 and fitting_enriched_X.shape[1] == 0:
+                    self.__log_warning(self.bundle.get("metrics_no_important_free_features"))
+                    return None
+
                 print(self.bundle.get("metrics_start"))
                 with Spinner():
-                    if fitting_X.shape[1] == 0 and fitting_enriched_X.shape[1] == 0:
-                        print(self.bundle.get("metrics_no_important_free_features"))
-                        self.logger.warning("No client or free relevant ADS features found to calculate metrics")
-                        self.warning_counter.increment()
-                        return None
-
                     self._check_train_and_eval_target_distribution(y_sorted, fitting_eval_set_dict)
 
-                    has_date = self._get_date_column(search_keys) is not None
+                    has_date = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME]) is not None
                     model_task_type = self.model_task_type or define_task(y_sorted, has_date, self.logger, silent=True)
 
                     wrapper = EstimatorWrapper.create(
@@ -930,11 +981,12 @@ class FeaturesEnricher(TransformerMixin):
                         scoring,
                         groups=groups,
                         text_features=self.generate_features,
+                        has_date=has_date,
                     )
                     metric = wrapper.metric_name
                     multiplier = wrapper.multiplier
 
-                    # 1 If client features are presented - fit and predict with KFold CatBoost model
+                    # 1 If client features are presented - fit and predict with KFold estimator
                     # on etalon features and calculate baseline metric
                     etalon_metric = None
                     baseline_estimator = None
@@ -956,14 +1008,24 @@ class FeaturesEnricher(TransformerMixin):
                             add_params=custom_loss_add_params,
                             groups=groups,
                             text_features=self.generate_features,
+                            has_date=has_date,
                         )
-                        etalon_metric = baseline_estimator.cross_val_predict(
+                        etalon_cv_result = baseline_estimator.cross_val_predict(
                             fitting_X, y_sorted, self.baseline_score_column
                         )
-                        self.logger.info(f"Baseline {metric} on train client features: {etalon_metric}")
+                        etalon_metric = etalon_cv_result.get_display_metric()
+                        if etalon_metric is None:
+                            self.logger.info(
+                                f"Baseline {metric} on train client features is None (maybe all features was removed)"
+                            )
+                            baseline_estimator = None
+                        else:
+                            self.logger.info(f"Baseline {metric} on train client features: {etalon_metric}")
 
-                    # 2 Fit and predict with KFold Catboost model on enriched tds
+                    # 2 Fit and predict with KFold estimator on enriched tds
                     # and calculate final metric (and uplift)
+                    enriched_metric = None
+                    uplift = None
                     enriched_estimator = None
                     if set(fitting_X.columns) != set(fitting_enriched_X.columns):
                         self.logger.info(
@@ -981,16 +1043,24 @@ class FeaturesEnricher(TransformerMixin):
                             add_params=custom_loss_add_params,
                             groups=groups,
                             text_features=self.generate_features,
+                            has_date=has_date,
                         )
-                        enriched_metric = enriched_estimator.cross_val_predict(fitting_enriched_X, enriched_y_sorted)
-                        self.logger.info(f"Enriched {metric} on train combined features: {enriched_metric}")
-                        if etalon_metric is not None:
-                            uplift = (enriched_metric - etalon_metric) * multiplier
+                        enriched_cv_result = enriched_estimator.cross_val_predict(fitting_enriched_X, enriched_y_sorted)
+                        enriched_metric = enriched_cv_result.get_display_metric()
+                        enriched_shaps = enriched_cv_result.shap_values
+
+                        if enriched_shaps is not None:
+                            self._update_shap_values(trace_id, validated_X.columns.to_list(), enriched_shaps)
+
+                        if enriched_metric is None:
+                            self.logger.warning(
+                                f"Enriched {metric} on train combined features is None (maybe all features was removed)"
+                            )
+                            enriched_estimator = None
                         else:
-                            uplift = None
-                    else:
-                        enriched_metric = None
-                        uplift = None
+                            self.logger.info(f"Enriched {metric} on train combined features: {enriched_metric}")
+                        if etalon_metric is not None and enriched_metric is not None:
+                            uplift = (enriched_cv_result.metric - etalon_cv_result.metric) * multiplier
 
                     train_metrics = {
                         self.bundle.get("quality_metrics_segment_header"): self.bundle.get(
@@ -999,10 +1069,10 @@ class FeaturesEnricher(TransformerMixin):
                         self.bundle.get("quality_metrics_rows_header"): _num_samples(effective_X),
                     }
                     if model_task_type in [ModelTaskType.BINARY, ModelTaskType.REGRESSION] and is_numeric_dtype(
-                        y_sorted
+                        validated_y
                     ):
                         train_metrics[self.bundle.get("quality_metrics_mean_target_header")] = round(
-                            np.mean(effective_y), 4
+                            np.mean(validated_y), 4
                         )
                     if etalon_metric is not None:
                         train_metrics[self.bundle.get("quality_metrics_baseline_header").format(metric)] = etalon_metric
@@ -1033,9 +1103,10 @@ class FeaturesEnricher(TransformerMixin):
                                     f"Calculate baseline {metric} on eval set {idx + 1} "
                                     f"on client features: {eval_X_sorted.columns.to_list()}"
                                 )
-                                etalon_eval_metric = baseline_estimator.calculate_metric(
+                                etalon_eval_results = baseline_estimator.calculate_metric(
                                     eval_X_sorted, eval_y_sorted, self.baseline_score_column
                                 )
+                                etalon_eval_metric = etalon_eval_results.get_display_metric()
                                 self.logger.info(
                                     f"Baseline {metric} on eval set {idx + 1} client features: {etalon_eval_metric}"
                                 )
@@ -1047,9 +1118,10 @@ class FeaturesEnricher(TransformerMixin):
                                     f"Calculate enriched {metric} on eval set {idx + 1} "
                                     f"on combined features: {enriched_eval_X_sorted.columns.to_list()}"
                                 )
-                                enriched_eval_metric = enriched_estimator.calculate_metric(
+                                enriched_eval_results = enriched_estimator.calculate_metric(
                                     enriched_eval_X_sorted, enriched_eval_y_sorted
                                 )
+                                enriched_eval_metric = enriched_eval_results.get_display_metric()
                                 self.logger.info(
                                     f"Enriched {metric} on eval set {idx + 1} combined features: {enriched_eval_metric}"
                                 )
@@ -1057,11 +1129,11 @@ class FeaturesEnricher(TransformerMixin):
                                 enriched_eval_metric = None
 
                             if etalon_eval_metric is not None and enriched_eval_metric is not None:
-                                eval_uplift = (enriched_eval_metric - etalon_eval_metric) * multiplier
+                                eval_uplift = (enriched_eval_results.metric - etalon_eval_results.metric) * multiplier
                             else:
                                 eval_uplift = None
 
-                            effective_eval_set = eval_set if eval_set is not None else self.eval_set
+                            # effective_eval_set = eval_set if eval_set is not None else self.eval_set
                             eval_metrics = {
                                 self.bundle.get("quality_metrics_segment_header"): self.bundle.get(
                                     "quality_metrics_eval_segment"
@@ -1072,10 +1144,10 @@ class FeaturesEnricher(TransformerMixin):
                                 # self.bundle.get("quality_metrics_match_rate_header"): eval_hit_rate,
                             }
                             if model_task_type in [ModelTaskType.BINARY, ModelTaskType.REGRESSION] and is_numeric_dtype(
-                                eval_y_sorted
+                                validated_eval_set[idx][1]
                             ):
                                 eval_metrics[self.bundle.get("quality_metrics_mean_target_header")] = round(
-                                    np.mean(effective_eval_set[idx][1]), 4
+                                    np.mean(validated_eval_set[idx][1]), 4
                                 )
                             if etalon_eval_metric is not None:
                                 eval_metrics[self.bundle.get("quality_metrics_baseline_header").format(metric)] = (
@@ -1099,7 +1171,7 @@ class FeaturesEnricher(TransformerMixin):
                     )
 
                     uplift_col = self.bundle.get("quality_metrics_uplift_header")
-                    date_column = self._get_date_column(search_keys)
+                    date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
                     if (
                         uplift_col in metrics_df.columns
                         and (metrics_df[uplift_col] < 0).any()
@@ -1138,6 +1210,57 @@ class FeaturesEnricher(TransformerMixin):
             finally:
                 self.logger.info(f"Calculating metrics elapsed time: {time.time() - start_time}")
 
+    def _update_shap_values(self, trace_id: str, x_columns: List[str], new_shaps: Dict[str, float]):
+        new_shaps = {
+            feature: _round_shap_value(shap) for feature, shap in new_shaps.items() if feature in self.feature_names_
+        }
+        self.__prepare_feature_importances(trace_id, x_columns, new_shaps, silent=True)
+
+        if self.features_info_display_handle is not None:
+            try:
+                _ = get_ipython()  # type: ignore
+
+                display_html_dataframe(
+                    self.features_info,
+                    self._features_info_without_links,
+                    self.bundle.get("relevant_features_header"),
+                    display_handle=self.features_info_display_handle,
+                )
+            except (ImportError, NameError):
+                pass
+        if self.data_sources_display_handle is not None:
+            try:
+                _ = get_ipython()  # type: ignore
+
+                display_html_dataframe(
+                    self.relevant_data_sources,
+                    self._relevant_data_sources_wo_links,
+                    self.bundle.get("relevant_data_sources_header"),
+                    display_handle=self.data_sources_display_handle,
+                )
+            except (ImportError, NameError):
+                pass
+        if self.autofe_features_display_handle is not None:
+            try:
+                _ = get_ipython()  # type: ignore
+                autofe_descriptions_df = self.get_autofe_features_description()
+                if autofe_descriptions_df is not None:
+                    display_html_dataframe(
+                        df=autofe_descriptions_df,
+                        internal_df=autofe_descriptions_df,
+                        header=self.bundle.get("autofe_descriptions_header"),
+                        display_handle=self.autofe_features_display_handle,
+                    )
+            except (ImportError, NameError):
+                pass
+        if self.report_button_handle is not None:
+            try:
+                _ = get_ipython()  # type: ignore
+
+                self.__show_report_button(display_handle=self.report_button_handle)
+            except (ImportError, NameError):
+                pass
+
     def _check_train_and_eval_target_distribution(self, y, eval_set_dict):
         uneven_distribution = False
         for eval_set in eval_set_dict.values():
@@ -1174,34 +1297,6 @@ class FeaturesEnricher(TransformerMixin):
     def _has_paid_features(self, exclude_features_sources: Optional[List[str]]) -> bool:
         return self._has_features_with_commercial_schema(CommercialSchema.PAID.value, exclude_features_sources)
 
-    def _extend_x(self, x: pd.DataFrame, is_demo_dataset: bool) -> Tuple[pd.DataFrame, Dict[str, SearchKey]]:
-        search_keys = self.search_keys.copy()
-        search_keys = self.__prepare_search_keys(x, search_keys, is_demo_dataset, is_transform=True, silent_mode=True)
-
-        extended_X = x.copy()
-        generated_features = []
-        date_column = self._get_date_column(search_keys)
-        if date_column is not None:
-            converter = DateTimeSearchKeyConverter(date_column, self.date_format, self.logger, self.bundle)
-            extended_X = converter.convert(extended_X, keep_time=True)
-            generated_features.extend(converter.generated_features)
-        email_column = self._get_email_column(search_keys)
-        hem_column = self._get_hem_column(search_keys)
-        if email_column:
-            converter = EmailSearchKeyConverter(email_column, hem_column, search_keys, self.logger)
-            extended_X = converter.convert(extended_X)
-            generated_features.extend(converter.generated_features)
-        if (
-            self.detect_missing_search_keys
-            and list(search_keys.values()) == [SearchKey.DATE]
-            and self.country_code is None
-        ):
-            converter = IpToCountrySearchKeyConverter(search_keys, self.logger)
-            extended_X = converter.convert(extended_X)
-        generated_features = [f for f in generated_features if f in self.fit_generated_features]
-
-        return extended_X, search_keys
-
     def _is_input_same_as_fit(
         self,
         X: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
@@ -1245,7 +1340,7 @@ class FeaturesEnricher(TransformerMixin):
         groups = None
 
         if not isinstance(_cv, BaseCrossValidator):
-            date_column = self._get_date_column(search_keys)
+            date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
             date_series = X[date_column] if date_column is not None else None
             _cv, groups = CVConfig(
                 _cv, date_series, self.random_state, self._search_task.get_shuffle_kfold(), group_columns=group_columns
@@ -1268,7 +1363,7 @@ class FeaturesEnricher(TransformerMixin):
 
     def _get_client_cat_features(
         self, estimator: Optional[Any], X: pd.DataFrame, search_keys: Dict[str, SearchKey]
-    ) -> Optional[List[str]]:
+    ) -> Tuple[Optional[List[str]], List[str]]:
         cat_features = None
         search_keys_for_metrics = []
         if (
@@ -1328,30 +1423,38 @@ class FeaturesEnricher(TransformerMixin):
             progress_bar,
             progress_callback,
         )
-        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys = dataclasses.astuple(sampled_data)
+        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming = dataclasses.astuple(
+            sampled_data
+        )
 
         excluding_search_keys = list(search_keys.keys())
         if search_keys_for_metrics is not None and len(search_keys_for_metrics) > 0:
-            excluding_search_keys = [sk for sk in excluding_search_keys if sk not in search_keys_for_metrics]
-        meta = self._search_task.get_all_features_metadata_v2()
-        zero_importance_client_features = [m for m in meta if m.source == "etalon" and m.shap_value == 0.0]
+            for sk in excluding_search_keys:
+                if columns_renaming.get(sk) in search_keys_for_metrics:
+                    excluding_search_keys.remove(sk)
 
         client_features = [
             c
             for c in X_sampled.columns.to_list()
-            if c
+            if (
+                not self.select_features
+                or c in self.feature_names_
+                or (self.fit_columns_renaming is not None and self.fit_columns_renaming.get(c) in self.feature_names_)
+            )
+            and c
             not in (
                 excluding_search_keys
                 + list(self.fit_dropped_features)
-                + [DateTimeSearchKeyConverter.DATETIME_COL, SYSTEM_RECORD_ID]
-                + zero_importance_client_features
+                + [DateTimeSearchKeyConverter.DATETIME_COL, SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID]
             )
         ]
+        self.logger.info(f"Client features column on prepare data for metrics: {client_features}")
 
         filtered_enriched_features = self.__filtered_enriched_features(
             importance_threshold,
             max_features,
         )
+        filtered_enriched_features = [c for c in filtered_enriched_features if c not in client_features]
 
         X_sorted, y_sorted = self._sort_by_system_record_id(X_sampled, y_sampled, self.cv)
         enriched_X_sorted, enriched_y_sorted = self._sort_by_system_record_id(enriched_X, y_sampled, self.cv)
@@ -1381,9 +1484,12 @@ class FeaturesEnricher(TransformerMixin):
             fitting_X = fitting_X.drop(columns=constant_columns, errors="ignore")
             fitting_enriched_X = fitting_enriched_X.drop(columns=constant_columns, errors="ignore")
 
+        # TODO maybe there is no more need for these convertions
         # Remove datetime features
         datetime_features = [
-            f for f in fitting_X.columns if is_datetime64_any_dtype(fitting_X[f]) or is_period_dtype(fitting_X[f])
+            f
+            for f in fitting_X.columns
+            if is_datetime64_any_dtype(fitting_X[f]) or isinstance(fitting_X[f].dtype, pd.PeriodDtype)
         ]
         if len(datetime_features) > 0:
             self.logger.warning(self.bundle.get("dataset_date_features").format(datetime_features))
@@ -1403,37 +1509,25 @@ class FeaturesEnricher(TransformerMixin):
         if len(decimal_columns_to_fix) > 0:
             self.logger.warning(f"Convert strings with decimal comma to float: {decimal_columns_to_fix}")
             for col in decimal_columns_to_fix:
-                fitting_X[col] = fitting_X[col].astype("string").str.replace(",", ".").astype(np.float64)
+                fitting_X[col] = fitting_X[col].astype("string").str.replace(",", ".", regex=False).astype(np.float64)
                 fitting_enriched_X[col] = (
-                    fitting_enriched_X[col].astype("string").str.replace(",", ".").astype(np.float64)
+                    fitting_enriched_X[col].astype("string").str.replace(",", ".", regex=False).astype(np.float64)
                 )
 
-        fitting_eval_set_dict = dict()
+        fitting_eval_set_dict = {}
+        fitting_x_columns = fitting_X.columns.to_list()
+        self.logger.info(f"Final list of fitting X columns: {fitting_x_columns}")
+        fitting_enriched_x_columns = fitting_enriched_X.columns.to_list()
+        self.logger.info(f"Final list of fitting enriched X columns: {fitting_enriched_x_columns}")
         for idx, eval_tuple in eval_set_sampled_dict.items():
             eval_X_sampled, enriched_eval_X, eval_y_sampled = eval_tuple
             eval_X_sorted, eval_y_sorted = self._sort_by_system_record_id(eval_X_sampled, eval_y_sampled, self.cv)
             enriched_eval_X_sorted, enriched_eval_y_sorted = self._sort_by_system_record_id(
                 enriched_eval_X, eval_y_sampled, self.cv
             )
-            fitting_eval_X = eval_X_sorted[client_features].copy()
-            fitting_enriched_eval_X = enriched_eval_X_sorted[
-                client_features + existing_filtered_enriched_features
-            ].copy()
-
-            # # Drop high cardinality features in eval set
-            if len(columns_with_high_cardinality) > 0:
-                fitting_eval_X = fitting_eval_X.drop(columns=columns_with_high_cardinality, errors="ignore")
-                fitting_enriched_eval_X = fitting_enriched_eval_X.drop(
-                    columns=columns_with_high_cardinality, errors="ignore"
-                )
-            # Drop constant features in eval_set
-            if len(constant_columns) > 0:
-                fitting_eval_X = fitting_eval_X.drop(columns=constant_columns, errors="ignore")
-                fitting_enriched_eval_X = fitting_enriched_eval_X.drop(columns=constant_columns, errors="ignore")
-            # Drop datetime features in eval_set
-            if len(datetime_features) > 0:
-                fitting_eval_X = fitting_eval_X.drop(columns=datetime_features, errors="ignore")
-                fitting_enriched_eval_X = fitting_enriched_eval_X.drop(columns=datetime_features, errors="ignore")
+            fitting_eval_X = eval_X_sorted[fitting_x_columns].copy()
+            fitting_enriched_eval_X = enriched_eval_X_sorted[fitting_enriched_x_columns].copy()
+
             # Convert bool to string in eval_set
             if len(bool_columns) > 0:
                 fitting_eval_X[col] = fitting_eval_X[col].astype(str)
@@ -1441,9 +1535,14 @@ class FeaturesEnricher(TransformerMixin):
             # Correct string features with decimal commas
             if len(decimal_columns_to_fix) > 0:
                 for col in decimal_columns_to_fix:
-                    fitting_eval_X[col] = fitting_eval_X[col].astype("string").str.replace(",", ".").astype(np.float64)
+                    fitting_eval_X[col] = (
+                        fitting_eval_X[col].astype("string").str.replace(",", ".", regex=False).astype(np.float64)
+                    )
                     fitting_enriched_eval_X[col] = (
-                        fitting_enriched_eval_X[col].astype("string").str.replace(",", ".").astype(np.float64)
+                        fitting_enriched_eval_X[col]
+                        .astype("string")
+                        .str.replace(",", ".", regex=False)
+                        .astype(np.float64)
                     )
 
             fitting_eval_set_dict[idx] = (
@@ -1463,6 +1562,7 @@ class FeaturesEnricher(TransformerMixin):
             search_keys,
             groups,
             cv,
+            columns_renaming,
         )
 
     @dataclass
@@ -1472,6 +1572,7 @@ class FeaturesEnricher(TransformerMixin):
         enriched_X: pd.DataFrame
         eval_set_sampled_dict: Dict[int, Tuple[pd.DataFrame, pd.Series]]
         search_keys: Dict[str, SearchKey]
+        columns_renaming: Dict[str, str]
 
     def _sample_data_for_metrics(
         self,
@@ -1486,18 +1587,28 @@ class FeaturesEnricher(TransformerMixin):
         progress_bar: Optional[ProgressBar],
         progress_callback: Optional[Callable[[SearchProgress], Any]],
     ) -> _SampledDataForMetrics:
-        if self.__cached_sampled_datasets is not None and is_input_same_as_fit and remove_outliers_calc_metrics is None:
+        datasets_hash = hash_input(validated_X, validated_y, eval_set)
+        cached_sampled_datasets = self.__cached_sampled_datasets.get(datasets_hash)
+        if cached_sampled_datasets is not None and is_input_same_as_fit and remove_outliers_calc_metrics is None:
             self.logger.info("Cached enriched dataset found - use it")
-            return self.__get_sampled_cached_enriched(exclude_features_sources)
+            return self.__get_sampled_cached_enriched(datasets_hash, exclude_features_sources)
         elif len(self.feature_importances_) == 0:
             self.logger.info("No external features selected. So use only input datasets for metrics calculation")
             return self.__sample_only_input(validated_X, validated_y, eval_set, is_demo_dataset)
         # TODO save and check if dataset was deduplicated - use imbalance branch for such case
-        elif not self.imbalanced and not exclude_features_sources and is_input_same_as_fit:
+        elif (
+            not self.imbalanced
+            and not exclude_features_sources
+            and is_input_same_as_fit
+            and self.df_with_original_index is not None
+        ):
             self.logger.info("Dataset is not imbalanced, so use enriched_X from fit")
             return self.__sample_balanced(eval_set, trace_id, remove_outliers_calc_metrics)
         else:
-            self.logger.info("Dataset is imbalanced or exclude_features_sources or X was passed. Run transform")
+            self.logger.info(
+                "Dataset is imbalanced or exclude_features_sources or X was passed or this is saved search."
+                " Run transform"
+            )
             print(self.bundle.get("prepare_data_for_metrics"))
             return self.__sample_imbalanced(
                 validated_X,
@@ -1510,17 +1621,23 @@ class FeaturesEnricher(TransformerMixin):
                 progress_callback,
             )
 
-    def __get_sampled_cached_enriched(self, exclude_features_sources: Optional[List[str]]) -> _SampledDataForMetrics:
-        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys = self.__cached_sampled_datasets
+    def __get_sampled_cached_enriched(
+        self, datasets_hash: str, exclude_features_sources: Optional[List[str]]
+    ) -> _SampledDataForMetrics:
+        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming = (
+            self.__cached_sampled_datasets[datasets_hash]
+        )
         if exclude_features_sources:
             enriched_X = enriched_X.drop(columns=exclude_features_sources, errors="ignore")
 
-        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+        return self.__mk_sampled_data_tuple(
+            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming
+        )
 
     def __sample_only_input(
         self, validated_X: pd.DataFrame, validated_y: pd.Series, eval_set: Optional[List[tuple]], is_demo_dataset: bool
     ) -> _SampledDataForMetrics:
-        eval_set_sampled_dict = dict()
+        eval_set_sampled_dict = {}
 
         df = validated_X.copy()
         df[TARGET] = validated_y
@@ -1533,7 +1650,31 @@ class FeaturesEnricher(TransformerMixin):
                 eval_xy[EVAL_SET_INDEX] = idx + 1
                 df = pd.concat([df, eval_xy])
 
-        df = clean_full_duplicates(df, logger=self.logger, silent=True, bundle=self.bundle)
+        search_keys = self.search_keys.copy()
+        search_keys = self.__prepare_search_keys(df, search_keys, is_demo_dataset, is_transform=True, silent_mode=True)
+
+        date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+        generated_features = []
+        if date_column is not None:
+            converter = DateTimeSearchKeyConverter(date_column, self.date_format, self.logger, self.bundle)
+            # Leave original date column values
+            df_with_date_features = converter.convert(df, keep_time=True)
+            df_with_date_features[date_column] = df[date_column]
+            df = df_with_date_features
+            generated_features = converter.generated_features
+
+        email_columns = SearchKey.find_all_keys(search_keys, SearchKey.EMAIL)
+        if email_columns:
+            generator = EmailDomainGenerator(email_columns)
+            df = generator.generate(df)
+            generated_features.extend(generator.generated_features)
+
+        # normalizer = Normalizer(self.bundle, self.logger)
+        # df, search_keys, generated_features = normalizer.normalize(df, search_keys, generated_features)
+        # columns_renaming = normalizer.columns_renaming
+        columns_renaming = {c: c for c in df.columns}
+
+        df, _ = clean_full_duplicates(df, logger=self.logger, bundle=self.bundle)
 
         num_samples = _num_samples(df)
         sample_threshold, sample_rows = (
@@ -1545,24 +1686,36 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.info(f"Downsampling from {num_samples} to {sample_rows}")
             df = df.sample(n=sample_rows, random_state=self.random_state)
 
-        df_extended, search_keys = self._extend_x(df, is_demo_dataset)
-        df_extended = self.__add_fit_system_record_id(df_extended, dict(), search_keys)
+        df = self.__add_fit_system_record_id(df, search_keys, SYSTEM_RECORD_ID)
+        if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
+            df = df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL)
 
-        train_df = df_extended.query(f"{EVAL_SET_INDEX} == 0") if eval_set is not None else df_extended
+        train_df = df.query(f"{EVAL_SET_INDEX} == 0") if eval_set is not None else df
         X_sampled = train_df.drop(columns=[TARGET, EVAL_SET_INDEX], errors="ignore")
         y_sampled = train_df[TARGET].copy()
         enriched_X = X_sampled
 
         if eval_set is not None:
             for idx in range(len(eval_set)):
-                eval_xy_sampled = df_extended.query(f"{EVAL_SET_INDEX} == {idx + 1}")
+                eval_xy_sampled = df.query(f"{EVAL_SET_INDEX} == {idx + 1}")
                 eval_X_sampled = eval_xy_sampled.drop(columns=[TARGET, EVAL_SET_INDEX], errors="ignore")
                 eval_y_sampled = eval_xy_sampled[TARGET].copy()
                 enriched_eval_X = eval_X_sampled
                 eval_set_sampled_dict[idx] = (eval_X_sampled, enriched_eval_X, eval_y_sampled)
-        self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
 
-        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+        datasets_hash = hash_input(X_sampled, y_sampled, eval_set_sampled_dict)
+        self.__cached_sampled_datasets[datasets_hash] = (
+            X_sampled,
+            y_sampled,
+            enriched_X,
+            eval_set_sampled_dict,
+            search_keys,
+            columns_renaming,
+        )
+
+        return self.__mk_sampled_data_tuple(
+            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming
+        )
 
     def __sample_balanced(
         self,
@@ -1570,22 +1723,21 @@ class FeaturesEnricher(TransformerMixin):
         trace_id: str,
         remove_outliers_calc_metrics: Optional[bool],
     ) -> _SampledDataForMetrics:
-        eval_set_sampled_dict = dict()
+        eval_set_sampled_dict = {}
         search_keys = self.fit_search_keys
 
         rows_to_drop = None
-        has_date = self._get_date_column(search_keys) is not None
-        task_type = self.model_task_type or define_task(
+        has_date = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME]) is not None
+        self.model_task_type = self.model_task_type or define_task(
             self.df_with_original_index[TARGET], has_date, self.logger, silent=True
         )
-        if task_type == ModelTaskType.REGRESSION:
+        if self.model_task_type == ModelTaskType.REGRESSION:
             target_outliers_df = self._search_task.get_target_outliers(trace_id)
             if target_outliers_df is not None and len(target_outliers_df) > 0:
                 outliers = pd.merge(
                     self.df_with_original_index,
                     target_outliers_df,
-                    left_on=SYSTEM_RECORD_ID,
-                    right_on=SYSTEM_RECORD_ID,
+                    on=ENTITY_SYSTEM_RECORD_ID,
                     how="inner",
                 )
                 top_outliers = outliers.sort_values(by=TARGET, ascending=False)[TARGET].head(3)
@@ -1612,6 +1764,7 @@ class FeaturesEnricher(TransformerMixin):
         X_sampled = enriched_Xy[x_columns].copy()
         y_sampled = enriched_Xy[TARGET].copy()
         enriched_X = enriched_Xy.drop(columns=[TARGET, EVAL_SET_INDEX], errors="ignore")
+        enriched_X_columns = enriched_X.columns.to_list()
 
         self.logger.info(f"Shape of enriched_X: {enriched_X.shape}")
         self.logger.info(f"Shape of X after sampling: {X_sampled.shape}")
@@ -1626,12 +1779,22 @@ class FeaturesEnricher(TransformerMixin):
             for idx in range(len(eval_set)):
                 eval_X_sampled = enriched_eval_sets[idx + 1][x_columns].copy()
                 eval_y_sampled = enriched_eval_sets[idx + 1][TARGET].copy()
-                enriched_eval_X = enriched_eval_sets[idx + 1].drop(columns=[TARGET, EVAL_SET_INDEX])
+                enriched_eval_X = enriched_eval_sets[idx + 1][enriched_X_columns].copy()
                 eval_set_sampled_dict[idx] = (eval_X_sampled, enriched_eval_X, eval_y_sampled)
 
-        self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+        datasets_hash = hash_input(self.X, self.y, self.eval_set)
+        self.__cached_sampled_datasets[datasets_hash] = (
+            X_sampled,
+            y_sampled,
+            enriched_X,
+            eval_set_sampled_dict,
+            search_keys,
+            self.fit_columns_renaming,
+        )
 
-        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+        return self.__mk_sampled_data_tuple(
+            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, self.fit_columns_renaming
+        )
 
     def __sample_imbalanced(
         self,
@@ -1644,7 +1807,7 @@ class FeaturesEnricher(TransformerMixin):
         progress_bar: Optional[ProgressBar],
         progress_callback: Optional[Callable[[SearchProgress], Any]],
     ) -> _SampledDataForMetrics:
-        eval_set_sampled_dict = dict()
+        eval_set_sampled_dict = {}
         if eval_set is not None:
             self.logger.info("Transform with eval_set")
             # concatenate X and eval_set with eval_set_index
@@ -1658,7 +1821,7 @@ class FeaturesEnricher(TransformerMixin):
                 eval_df_with_index[EVAL_SET_INDEX] = idx + 1
                 df = pd.concat([df, eval_df_with_index])
 
-            df = clean_full_duplicates(df, logger=self.logger, silent=True, bundle=self.bundle)
+            df, _ = clean_full_duplicates(df, logger=self.logger, bundle=self.bundle)
 
             # downsample if need to eval_set threshold
             num_samples = _num_samples(df)
@@ -1666,12 +1829,12 @@ class FeaturesEnricher(TransformerMixin):
                 self.logger.info(f"Downsampling from {num_samples} to {Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS}")
                 df = df.sample(n=Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS, random_state=self.random_state)
 
-            eval_set_sampled_dict = dict()
+            eval_set_sampled_dict = {}
 
             tmp_target_name = "__target"
             df = df.rename(columns={TARGET: tmp_target_name})
 
-            enriched_df = self.__inner_transform(
+            enriched_df, columns_renaming, generated_features = self.__inner_transform(
                 trace_id,
                 df,
                 exclude_features_sources=exclude_features_sources,
@@ -1688,7 +1851,7 @@ class FeaturesEnricher(TransformerMixin):
 
             x_columns = [
                 c
-                for c in (validated_X.columns.tolist() + self.fit_generated_features + [SYSTEM_RECORD_ID])
+                for c in (validated_X.columns.tolist() + generated_features + [SYSTEM_RECORD_ID])
                 if c in enriched_df.columns
             ]
 
@@ -1696,12 +1859,13 @@ class FeaturesEnricher(TransformerMixin):
             X_sampled = enriched_Xy[x_columns].copy()
             y_sampled = enriched_Xy[TARGET].copy()
             enriched_X = enriched_Xy.drop(columns=[TARGET, EVAL_SET_INDEX])
+            enriched_X_columns = enriched_X.columns.tolist()
 
             for idx in range(len(eval_set)):
                 enriched_eval_xy = enriched_df.query(f"{EVAL_SET_INDEX} == {idx + 1}")
                 eval_x_sampled = enriched_eval_xy[x_columns].copy()
                 eval_y_sampled = enriched_eval_xy[TARGET].copy()
-                enriched_eval_x = enriched_eval_xy.drop(columns=[TARGET, EVAL_SET_INDEX])
+                enriched_eval_x = enriched_eval_xy[enriched_X_columns].copy()
                 eval_set_sampled_dict[idx] = (eval_x_sampled, enriched_eval_x, eval_y_sampled)
         else:
             self.logger.info("Transform without eval_set")
@@ -1709,7 +1873,7 @@ class FeaturesEnricher(TransformerMixin):
 
             df[TARGET] = validated_y
 
-            df = clean_full_duplicates(df, logger=self.logger, silent=True, bundle=self.bundle)
+            df, _ = clean_full_duplicates(df, logger=self.logger, bundle=self.bundle)
 
             num_samples = _num_samples(df)
             if num_samples > Dataset.FIT_SAMPLE_THRESHOLD:
@@ -1719,7 +1883,7 @@ class FeaturesEnricher(TransformerMixin):
             tmp_target_name = "__target"
             df = df.rename(columns={TARGET: tmp_target_name})
 
-            enriched_Xy = self.__inner_transform(
+            enriched_Xy, columns_renaming, generated_features = self.__inner_transform(
                 trace_id,
                 df,
                 exclude_features_sources=exclude_features_sources,
@@ -1736,7 +1900,7 @@ class FeaturesEnricher(TransformerMixin):
 
             x_columns = [
                 c
-                for c in (validated_X.columns.tolist() + self.fit_generated_features + [SYSTEM_RECORD_ID])
+                for c in (validated_X.columns.tolist() + generated_features + [SYSTEM_RECORD_ID])
                 if c in enriched_Xy.columns
             ]
 
@@ -1744,9 +1908,19 @@ class FeaturesEnricher(TransformerMixin):
             y_sampled = enriched_Xy[TARGET].copy()
             enriched_X = enriched_Xy.drop(columns=TARGET)
 
-        self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, self.search_keys)
+        datasets_hash = hash_input(validated_X, validated_y, eval_set)
+        self.__cached_sampled_datasets[datasets_hash] = (
+            X_sampled,
+            y_sampled,
+            enriched_X,
+            eval_set_sampled_dict,
+            self.search_keys,
+            columns_renaming,
+        )
 
-        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, self.search_keys)
+        return self.__mk_sampled_data_tuple(
+            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, self.search_keys, columns_renaming
+        )
 
     def __mk_sampled_data_tuple(
         self,
@@ -1755,6 +1929,7 @@ class FeaturesEnricher(TransformerMixin):
         enriched_X: pd.DataFrame,
         eval_set_sampled_dict: Dict,
         search_keys: Dict,
+        columns_renaming: Dict[str, str],
     ):
         search_keys = {k: v for k, v in search_keys.items() if k in X_sampled.columns.to_list()}
         return FeaturesEnricher._SampledDataForMetrics(
@@ -1763,6 +1938,7 @@ class FeaturesEnricher(TransformerMixin):
             enriched_X=enriched_X,
             eval_set_sampled_dict=eval_set_sampled_dict,
             search_keys=search_keys,
+            columns_renaming=columns_renaming,
         )
 
     def get_search_id(self) -> Optional[str]:
@@ -1812,9 +1988,19 @@ class FeaturesEnricher(TransformerMixin):
         file_metadata = self._search_task.get_file_metadata(str(uuid.uuid4()))
         search_keys = file_metadata.search_types()
         if SearchKey.IPV6_ADDRESS in search_keys:
-            search_keys.remove(SearchKey.IPV6_ADDRESS)
+            # search_keys.remove(SearchKey.IPV6_ADDRESS)
+            search_keys.pop(SearchKey.IPV6_ADDRESS, None)
 
-        keys = "{" + ", ".join([f'"{key.name}": "{key_example(key)}"' for key in search_keys]) + "}"
+        keys = (
+            "{"
+            + ", ".join(
+                [
+                    f'"{key.name}": {{"name": "{name}", "value": "{key_example(key)}"}}'
+                    for key, name in search_keys.items()
+                ]
+            )
+            + "}"
+        )
         features_for_transform = self._search_task.get_features_for_transform()
         if features_for_transform:
             original_features_for_transform = [
@@ -1851,37 +2037,41 @@ class FeaturesEnricher(TransformerMixin):
         progress_bar: Optional[ProgressBar] = None,
         progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
         add_fit_system_record_id: bool = False,
-    ) -> pd.DataFrame:
+    ) -> Tuple[pd.DataFrame, Dict[str, str], List[str]]:
         if self._search_task is None:
             raise NotFittedError(self.bundle.get("transform_unfitted_enricher"))
 
         start_time = time.time()
         with MDC(trace_id=trace_id):
             self.logger.info("Start transform")
-            self.__log_debug_information(X, exclude_features_sources=exclude_features_sources)
+
+            validated_X = self._validate_X(X, is_transform=True)
+
+            self.__log_debug_information(validated_X, exclude_features_sources=exclude_features_sources)
 
             self.__validate_search_keys(self.search_keys, self.search_id)
 
             if len(self.feature_names_) == 0:
                 self.logger.warning(self.bundle.get("no_important_features_for_transform"))
-                return X
+                return X, {c: c for c in X.columns}, []
 
             if self._has_paid_features(exclude_features_sources):
                 msg = self.bundle.get("transform_with_paid_features")
                 self.logger.warning(msg)
                 self.__display_support_link(msg)
-                return None
+                return None, {c: c for c in X.columns}, []
 
             if not metrics_calculation:
                 transform_usage = self.rest_client.get_current_transform_usage(trace_id)
                 self.logger.info(f"Current transform usage: {transform_usage}. Transforming {len(X)} rows")
                 if transform_usage.has_limit:
                     if len(X) > transform_usage.rest_rows:
-                        msg = self.bundle.get("transform_usage_warning").format(len(X), transform_usage.rest_rows)
+                        rest_rows = max(transform_usage.rest_rows, 0)
+                        msg = self.bundle.get("transform_usage_warning").format(len(X), rest_rows)
                         self.logger.warning(msg)
                         print(msg)
                         show_request_quote_button()
-                        return None
+                        return None, {c: c for c in X.columns}, []
                     else:
                         msg = self.bundle.get("transform_usage_info").format(
                             transform_usage.limit, transform_usage.transformed_rows
@@ -1889,11 +2079,11 @@ class FeaturesEnricher(TransformerMixin):
                         self.logger.info(msg)
                         print(msg)
 
-            validated_X = self._validate_X(X, is_transform=True)
-
             is_demo_dataset = hash_input(validated_X) in DEMO_DATASET_HASHES
 
-            columns_to_drop = [c for c in validated_X.columns if c in self.feature_names_]
+            columns_to_drop = [
+                c for c in validated_X.columns if c in self.feature_names_ and c in self.dropped_client_feature_names_
+            ]
             if len(columns_to_drop) > 0:
                 msg = self.bundle.get("x_contains_enriching_columns").format(columns_to_drop)
                 self.logger.warning(msg)
@@ -1919,79 +2109,135 @@ class FeaturesEnricher(TransformerMixin):
             df = self.__add_country_code(df, search_keys)
 
             generated_features = []
-            date_column = self._get_date_column(search_keys)
+            date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
             if date_column is not None:
                 converter = DateTimeSearchKeyConverter(date_column, self.date_format, self.logger, bundle=self.bundle)
-                df = converter.convert(df)
+                df = converter.convert(df, keep_time=True)
                 self.logger.info(f"Date column after convertion: {df[date_column]}")
                 generated_features.extend(converter.generated_features)
             else:
                 self.logger.info("Input dataset hasn't date column")
                 if self.add_date_if_missing:
                     df = self._add_current_date_as_key(df, search_keys, self.logger, self.bundle)
+
+            email_columns = SearchKey.find_all_keys(search_keys, SearchKey.EMAIL)
+            if email_columns:
+                generator = EmailDomainGenerator(email_columns)
+                df = generator.generate(df)
+                generated_features.extend(generator.generated_features)
+
+            normalizer = Normalizer(self.bundle, self.logger)
+            df, search_keys, generated_features = normalizer.normalize(df, search_keys, generated_features)
+            columns_renaming = normalizer.columns_renaming
+
+            # Don't pass all features in backend on transform
+            runtime_parameters = self._get_copy_of_runtime_parameters()
+            features_for_transform = self._search_task.get_features_for_transform() or []
+            if len(features_for_transform) > 0:
+                missing_features_for_transform = [
+                    columns_renaming.get(f) for f in features_for_transform if f not in df.columns
+                ]
+                if len(missing_features_for_transform) > 0:
+                    raise ValidationError(
+                        self.bundle.get("missing_features_for_transform").format(missing_features_for_transform)
+                    )
+                runtime_parameters.properties["features_for_embeddings"] = ",".join(features_for_transform)
+
+            columns_for_system_record_id = sorted(list(search_keys.keys()) + features_for_transform)
+
+            df[ENTITY_SYSTEM_RECORD_ID] = pd.util.hash_pandas_object(
+                df[columns_for_system_record_id], index=False
+            ).astype("float64")
+
+            # Explode multiple search keys
+            df, unnest_search_keys = self._explode_multiple_search_keys(df, search_keys, columns_renaming)
+
             email_column = self._get_email_column(search_keys)
             hem_column = self._get_hem_column(search_keys)
-            email_converted_to_hem = False
             if email_column:
-                converter = EmailSearchKeyConverter(email_column, hem_column, search_keys, self.logger)
+                converter = EmailSearchKeyConverter(
+                    email_column,
+                    hem_column,
+                    search_keys,
+                    columns_renaming,
+                    list(unnest_search_keys.keys()),
+                    self.logger,
+                )
                 df = converter.convert(df)
-                generated_features.extend(converter.generated_features)
-                email_converted_to_hem = converter.email_converted_to_hem
-            if (
-                self.detect_missing_search_keys
-                and list(search_keys.values()) == [SearchKey.DATE]
-                and self.country_code is None
-            ):
-                converter = IpToCountrySearchKeyConverter(search_keys, self.logger)
+
+            ip_column = self._get_ip_column(search_keys)
+            if ip_column:
+                converter = IpSearchKeyConverter(
+                    ip_column,
+                    search_keys,
+                    columns_renaming,
+                    list(unnest_search_keys.keys()),
+                    self.bundle,
+                    self.logger,
+                )
                 df = converter.convert(df)
-            generated_features = [f for f in generated_features if f in self.fit_generated_features]
 
-            meaning_types = {col: key.value for col, key in search_keys.items()}
-            non_keys_columns = [column for column in df.columns if column not in search_keys.keys()]
+            phone_column = self._get_phone_column(search_keys)
+            country_column = self._get_country_column(search_keys)
+            if phone_column:
+                converter = PhoneSearchKeyConverter(phone_column, country_column)
+                df = converter.convert(df)
+
+            if country_column:
+                converter = CountrySearchKeyConverter(country_column)
+                df = converter.convert(df)
 
-            if email_converted_to_hem:
-                non_keys_columns.append(email_column)
+            postal_code = self._get_postal_column(search_keys)
+            if postal_code:
+                converter = PostalCodeSearchKeyConverter(postal_code)
+                df = converter.convert(df)
 
-            # Don't pass features in backend on transform
-            original_features_for_transform = None
-            runtime_parameters = self._get_copy_of_runtime_parameters()
-            if len(non_keys_columns) > 0:
-                # Pass only features that need for transform
-                features_for_transform = self._search_task.get_features_for_transform()
-                if features_for_transform is not None and len(features_for_transform) > 0:
-                    file_metadata = self._search_task.get_file_metadata(trace_id)
-                    original_features_for_transform = [
-                        c.originalName or c.name for c in file_metadata.columns if c.name in features_for_transform
-                    ]
-                    non_keys_columns = [c for c in non_keys_columns if c not in original_features_for_transform]
+            # generated_features = [f for f in generated_features if f in self.fit_generated_features]
 
-                    runtime_parameters.properties["features_for_embeddings"] = ",".join(features_for_transform)
+            meaning_types = {col: key.value for col, key in search_keys.items()}
+            for col in features_for_transform:
+                meaning_types[col] = FileColumnMeaningType.FEATURE
+            features_not_to_pass = [
+                c
+                for c in df.columns
+                if c not in search_keys.keys()
+                and c not in features_for_transform
+                and c not in [ENTITY_SYSTEM_RECORD_ID, SEARCH_KEY_UNNEST]
+            ]
 
             if add_fit_system_record_id:
-                df = self.__add_fit_system_record_id(df, dict(), search_keys)
+                df = self.__add_fit_system_record_id(df, search_keys, SYSTEM_RECORD_ID)
                 df = df.rename(columns={SYSTEM_RECORD_ID: SORT_ID})
-                non_keys_columns.append(SORT_ID)
+                features_not_to_pass.append(SORT_ID)
 
-            columns_for_system_record_id = sorted(list(search_keys.keys()) + (original_features_for_transform or []))
+            if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
+                df = df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL)
 
+            # search keys might be changed after explode
+            columns_for_system_record_id = sorted(list(search_keys.keys()) + features_for_transform)
             df[SYSTEM_RECORD_ID] = pd.util.hash_pandas_object(df[columns_for_system_record_id], index=False).astype(
-                "Float64"
+                "float64"
             )
             meaning_types[SYSTEM_RECORD_ID] = FileColumnMeaningType.SYSTEM_RECORD_ID
+            meaning_types[ENTITY_SYSTEM_RECORD_ID] = FileColumnMeaningType.ENTITY_SYSTEM_RECORD_ID
+            if SEARCH_KEY_UNNEST in df.columns:
+                meaning_types[SEARCH_KEY_UNNEST] = FileColumnMeaningType.UNNEST_KEY
 
             df = df.reset_index(drop=True)
-            system_columns_with_original_index = [SYSTEM_RECORD_ID] + generated_features
+            system_columns_with_original_index = [SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID] + generated_features
             if add_fit_system_record_id:
                 system_columns_with_original_index.append(SORT_ID)
             df_with_original_index = df[system_columns_with_original_index].copy()
 
             combined_search_keys = combine_search_keys(search_keys.keys())
 
-            df_without_features = df.drop(columns=non_keys_columns)
+            df_without_features = df.drop(columns=features_not_to_pass, errors="ignore")
 
-            df_without_features = clean_full_duplicates(
-                df_without_features, self.logger, silent=silent_mode, bundle=self.bundle
+            df_without_features, full_duplicates_warning = clean_full_duplicates(
+                df_without_features, self.logger, bundle=self.bundle
             )
+            if not silent_mode and full_duplicates_warning:
+                self.__log_warning(full_duplicates_warning)
 
             del df
             gc.collect()
@@ -1999,14 +2245,14 @@ class FeaturesEnricher(TransformerMixin):
             dataset = Dataset(
                 "sample_" + str(uuid.uuid4()),
                 df=df_without_features,
+                meaning_types=meaning_types,
+                search_keys=combined_search_keys,
+                unnest_search_keys=unnest_search_keys,
                 date_format=self.date_format,
                 rest_client=self.rest_client,
                 logger=self.logger,
             )
-            dataset.meaning_types = meaning_types
-            dataset.search_keys = combined_search_keys
-            if email_converted_to_hem:
-                dataset.ignore_columns = [email_column]
+            dataset.columns_renaming = columns_renaming
 
             if max_features is not None or importance_threshold is not None:
                 exclude_features_sources = list(
@@ -2094,9 +2340,15 @@ class FeaturesEnricher(TransformerMixin):
             else:
                 result = enrich()
 
+            selecting_columns = [
+                c
+                for c in itertools.chain(validated_X.columns.tolist(), generated_features)
+                if c not in self.dropped_client_feature_names_
+            ]
             filtered_columns = self.__filtered_enriched_features(importance_threshold, max_features)
-            existing_filtered_columns = [c for c in filtered_columns if c in result.columns]
-            selecting_columns = validated_X.columns.tolist() + generated_features + existing_filtered_columns
+            selecting_columns.extend(
+                c for c in filtered_columns if c in result.columns and c not in validated_X.columns
+            )
             if add_fit_system_record_id:
                 selecting_columns.append(SORT_ID)
 
@@ -2108,7 +2360,7 @@ class FeaturesEnricher(TransformerMixin):
             if add_fit_system_record_id:
                 result = result.rename(columns={SORT_ID: SYSTEM_RECORD_ID})
 
-            return result
+            return result, columns_renaming, generated_features
 
     def _get_excluded_features(self, max_features: Optional[int], importance_threshold: Optional[float]) -> List[str]:
         features_info = self._internal_features_info
@@ -2132,7 +2384,7 @@ class FeaturesEnricher(TransformerMixin):
             ]
             return excluded_features[feature_name_header].values.tolist()
 
-    def __validate_search_keys(self, search_keys: Dict[str, SearchKey], search_id: Optional[str]):
+    def __validate_search_keys(self, search_keys: Dict[str, SearchKey], search_id: Optional[str] = None):
         if (search_keys is None or len(search_keys) == 0) and self.country_code is None:
             if search_id:
                 self.logger.debug(f"search_id {search_id} provided without search_keys")
@@ -2143,6 +2395,14 @@ class FeaturesEnricher(TransformerMixin):
 
         key_types = search_keys.values()
 
+        # Multiple search keys allowed only for PHONE, IP, POSTAL_CODE, EMAIL, HEM
+        multi_keys = [key for key, count in Counter(key_types).items() if count > 1]
+        for multi_key in multi_keys:
+            if multi_key not in [SearchKey.PHONE, SearchKey.IP, SearchKey.POSTAL_CODE, SearchKey.EMAIL, SearchKey.HEM]:
+                msg = self.bundle.get("unsupported_multi_key").format(multi_key)
+                self.logger.warning(msg)
+                raise ValidationError(msg)
+
         if SearchKey.DATE in key_types and SearchKey.DATETIME in key_types:
             msg = self.bundle.get("date_and_datetime_simultanious")
             self.logger.warning(msg)
@@ -2158,11 +2418,11 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(msg)
             raise ValidationError(msg)
 
-        for key_type in SearchKey.__members__.values():
-            if key_type != SearchKey.CUSTOM_KEY and list(key_types).count(key_type) > 1:
-                msg = self.bundle.get("multiple_search_key").format(key_type)
-                self.logger.warning(msg)
-                raise ValidationError(msg)
+        # for key_type in SearchKey.__members__.values():
+        #     if key_type != SearchKey.CUSTOM_KEY and list(key_types).count(key_type) > 1:
+        #         msg = self.bundle.get("multiple_search_key").format(key_type)
+        #         self.logger.warning(msg)
+        #         raise ValidationError(msg)
 
         # non_personal_keys = set(SearchKey.__members__.values()) - set(SearchKey.personal_keys())
         # if (
@@ -2178,6 +2438,15 @@ class FeaturesEnricher(TransformerMixin):
     def __is_registered(self) -> bool:
         return self.api_key is not None and self.api_key != ""
 
+    def __log_warning(self, message: str, show_support_link: bool = False):
+        warning_num = self.warning_counter.increment()
+        formatted_message = f"WARNING #{warning_num}: {message}\n"
+        if show_support_link:
+            self.__display_support_link(formatted_message)
+        else:
+            print(formatted_message)
+        self.logger.warning(message)
+
     def __inner_fit(
         self,
         trace_id: str,
@@ -2199,8 +2468,11 @@ class FeaturesEnricher(TransformerMixin):
     ):
         self.warning_counter.reset()
         self.df_with_original_index = None
-        self.__cached_sampled_datasets = None
+        self.__cached_sampled_datasets = dict()
         self.metrics = None
+        self.fit_columns_renaming = None
+        self.fit_dropped_features = set()
+        self.fit_generated_features = []
 
         validated_X = self._validate_X(X)
         validated_y = self._validate_y(validated_X, y)
@@ -2221,9 +2493,7 @@ class FeaturesEnricher(TransformerMixin):
             checked_generate_features = []
             for gen_feature in self.generate_features:
                 if gen_feature not in x_columns:
-                    msg = self.bundle.get("missing_generate_feature").format(gen_feature, x_columns)
-                    print(msg)
-                    self.logger.warning(msg)
+                    self.__log_warning(self.bundle.get("missing_generate_feature").format(gen_feature, x_columns))
                 else:
                     checked_generate_features.append(gen_feature)
             self.generate_features = checked_generate_features
@@ -2232,9 +2502,9 @@ class FeaturesEnricher(TransformerMixin):
         validate_scoring_argument(scoring)
 
         self.__log_debug_information(
-            X,
-            y,
-            eval_set,
+            validated_X,
+            validated_y,
+            validated_eval_set,
             exclude_features_sources=exclude_features_sources,
             calculate_metrics=calculate_metrics,
             scoring=scoring,
@@ -2244,20 +2514,6 @@ class FeaturesEnricher(TransformerMixin):
 
         df = pd.concat([validated_X, validated_y], axis=1)
 
-        self.fit_search_keys = self.search_keys.copy()
-        self.fit_search_keys = self.__prepare_search_keys(validated_X, self.fit_search_keys, is_demo_dataset)
-
-        validate_dates_distribution(validated_X, self.fit_search_keys, self.logger, self.bundle, self.warning_counter)
-
-        maybe_date_column = self._get_date_column(self.fit_search_keys)
-        has_date = maybe_date_column is not None
-        model_task_type = self.model_task_type or define_task(validated_y, has_date, self.logger)
-        self._validate_binary_observations(validated_y, model_task_type)
-
-        self.runtime_parameters = get_runtime_params_custom_loss(
-            self.loss, model_task_type, self.runtime_parameters, self.logger
-        )
-
         if validated_eval_set is not None and len(validated_eval_set) > 0:
             df[EVAL_SET_INDEX] = 0
             for idx, (eval_X, eval_y) in enumerate(validated_eval_set):
@@ -2265,12 +2521,21 @@ class FeaturesEnricher(TransformerMixin):
                 eval_df[EVAL_SET_INDEX] = idx + 1
                 df = pd.concat([df, eval_df])
 
-        df = self.__correct_target(df)
-
+        self.fit_search_keys = self.search_keys.copy()
         df = self.__handle_index_search_keys(df, self.fit_search_keys)
+        self.fit_search_keys = self.__prepare_search_keys(df, self.fit_search_keys, is_demo_dataset)
 
-        if is_numeric_dtype(df[self.TARGET_NAME]) and has_date:
-            self._validate_PSI(df.sort_values(by=maybe_date_column))
+        maybe_date_column = SearchKey.find_key(self.fit_search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+        has_date = maybe_date_column is not None
+        self.model_task_type = self.model_task_type or define_task(validated_y, has_date, self.logger)
+
+        self._validate_binary_observations(validated_y, self.model_task_type)
+
+        self.runtime_parameters = get_runtime_params_custom_loss(
+            self.loss, self.model_task_type, self.runtime_parameters, self.logger
+        )
+
+        df = self.__correct_target(df)
 
         if DEFAULT_INDEX in df.columns:
             msg = self.bundle.get("unsupported_index_column")
@@ -2281,58 +2546,132 @@ class FeaturesEnricher(TransformerMixin):
 
         df = self.__add_country_code(df, self.fit_search_keys)
 
-        df = remove_fintech_duplicates(
-            df, self.fit_search_keys, date_format=self.date_format, logger=self.logger, bundle=self.bundle
-        )
-        df = clean_full_duplicates(df, self.logger, bundle=self.bundle)
-
-        date_column = self._get_date_column(self.fit_search_keys)
-        self.__adjust_cv(df, date_column, model_task_type)
-
         self.fit_generated_features = []
 
-        if date_column is not None:
-            converter = DateTimeSearchKeyConverter(date_column, self.date_format, self.logger, bundle=self.bundle)
+        if has_date:
+            converter = DateTimeSearchKeyConverter(
+                maybe_date_column,
+                self.date_format,
+                self.logger,
+                bundle=self.bundle,
+            )
             df = converter.convert(df, keep_time=True)
-            self.logger.info(f"Date column after convertion: {df[date_column]}")
+            if converter.has_old_dates:
+                self.__log_warning(self.bundle.get("dataset_drop_old_dates"))
+            self.logger.info(f"Date column after convertion: {df[maybe_date_column]}")
             self.fit_generated_features.extend(converter.generated_features)
         else:
             self.logger.info("Input dataset hasn't date column")
             if self.add_date_if_missing:
                 df = self._add_current_date_as_key(df, self.fit_search_keys, self.logger, self.bundle)
+
+        email_columns = SearchKey.find_all_keys(self.fit_search_keys, SearchKey.EMAIL)
+        if email_columns:
+            generator = EmailDomainGenerator(email_columns)
+            df = generator.generate(df)
+            self.fit_generated_features.extend(generator.generated_features)
+
+        # Checks that need validated date
+        try:
+            if not is_dates_distribution_valid(df, self.fit_search_keys):
+                self.__log_warning(bundle.get("x_unstable_by_date"))
+        except Exception:
+            self.logger.exception("Failed to check dates distribution validity")
+
+        if (
+            is_numeric_dtype(df[self.TARGET_NAME])
+            and self.model_task_type in [ModelTaskType.BINARY, ModelTaskType.MULTICLASS]
+            and has_date
+        ):
+            self._validate_PSI(df.sort_values(by=maybe_date_column))
+
+        normalizer = Normalizer(self.bundle, self.logger)
+        df, self.fit_search_keys, self.fit_generated_features = normalizer.normalize(
+            df, self.fit_search_keys, self.fit_generated_features
+        )
+        self.fit_columns_renaming = normalizer.columns_renaming
+        if normalizer.removed_features:
+            self.__log_warning(self.bundle.get("dataset_date_features").format(normalizer.removed_features))
+
+        self.__adjust_cv(df)
+
+        df, fintech_warnings = remove_fintech_duplicates(
+            df, self.fit_search_keys, date_format=self.date_format, logger=self.logger, bundle=self.bundle
+        )
+        if fintech_warnings:
+            for fintech_warning in fintech_warnings:
+                self.__log_warning(fintech_warning)
+        df, full_duplicates_warning = clean_full_duplicates(df, self.logger, bundle=self.bundle)
+        if full_duplicates_warning:
+            self.__log_warning(full_duplicates_warning)
+
+        # Explode multiple search keys
+        df = self.__add_fit_system_record_id(df, self.fit_search_keys, ENTITY_SYSTEM_RECORD_ID)
+
+        # TODO check that this is correct for enrichment
+        self.df_with_original_index = df.copy()
+        # TODO check maybe need to drop _time column from df_with_original_index
+
+        df, unnest_search_keys = self._explode_multiple_search_keys(df, self.fit_search_keys, self.fit_columns_renaming)
+
+        # Convert EMAIL to HEM after unnesting to do it only with one column
         email_column = self._get_email_column(self.fit_search_keys)
         hem_column = self._get_hem_column(self.fit_search_keys)
-        email_converted_to_hem = False
         if email_column:
-            converter = EmailSearchKeyConverter(email_column, hem_column, self.fit_search_keys, self.logger)
+            converter = EmailSearchKeyConverter(
+                email_column,
+                hem_column,
+                self.fit_search_keys,
+                self.fit_columns_renaming,
+                list(unnest_search_keys.keys()),
+                self.logger,
+            )
             df = converter.convert(df)
-            self.fit_generated_features.extend(converter.generated_features)
-            email_converted_to_hem = converter.email_converted_to_hem
-        if (
-            self.detect_missing_search_keys
-            and list(self.fit_search_keys.values()) == [SearchKey.DATE]
-            and self.country_code is None
-        ):
-            converter = IpToCountrySearchKeyConverter(self.fit_search_keys, self.logger)
+
+        ip_column = self._get_ip_column(self.fit_search_keys)
+        if ip_column:
+            converter = IpSearchKeyConverter(
+                ip_column,
+                self.fit_search_keys,
+                self.fit_columns_renaming,
+                list(unnest_search_keys.keys()),
+                self.bundle,
+                self.logger,
+            )
+            df = converter.convert(df)
+
+        phone_column = self._get_phone_column(self.fit_search_keys)
+        country_column = self._get_country_column(self.fit_search_keys)
+        if phone_column:
+            converter = PhoneSearchKeyConverter(phone_column, country_column)
+            df = converter.convert(df)
+
+        if country_column:
+            converter = CountrySearchKeyConverter(country_column)
             df = converter.convert(df)
 
-        non_feature_columns = [self.TARGET_NAME, EVAL_SET_INDEX] + list(self.fit_search_keys.keys())
-        if email_converted_to_hem:
-            non_feature_columns.append(email_column)
+        postal_code = self._get_postal_column(self.fit_search_keys)
+        if postal_code:
+            converter = PostalCodeSearchKeyConverter(postal_code)
+            df = converter.convert(df)
+
+        non_feature_columns = [self.TARGET_NAME, EVAL_SET_INDEX, ENTITY_SYSTEM_RECORD_ID, SEARCH_KEY_UNNEST] + list(
+            self.fit_search_keys.keys()
+        )
         if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
             non_feature_columns.append(DateTimeSearchKeyConverter.DATETIME_COL)
 
         features_columns = [c for c in df.columns if c not in non_feature_columns]
 
-        features_to_drop = FeaturesValidator(self.logger).validate(
-            df, features_columns, self.generate_features, self.warning_counter
+        features_to_drop, feature_validator_warnings = FeaturesValidator(self.logger).validate(
+            df, features_columns, self.generate_features, self.fit_columns_renaming
         )
+        if feature_validator_warnings:
+            for warning in feature_validator_warnings:
+                self.__log_warning(warning)
         self.fit_dropped_features.update(features_to_drop)
         df = df.drop(columns=features_to_drop)
 
-        if email_converted_to_hem:
-            self.fit_dropped_features.add(email_column)
-
         self.fit_generated_features = [f for f in self.fit_generated_features if f not in self.fit_dropped_features]
 
         meaning_types = {
@@ -2340,12 +2679,19 @@ class FeaturesEnricher(TransformerMixin):
             **{str(c): FileColumnMeaningType.FEATURE for c in df.columns if c not in non_feature_columns},
         }
         meaning_types[self.TARGET_NAME] = FileColumnMeaningType.TARGET
+        meaning_types[ENTITY_SYSTEM_RECORD_ID] = FileColumnMeaningType.ENTITY_SYSTEM_RECORD_ID
+        if SEARCH_KEY_UNNEST in df.columns:
+            meaning_types[SEARCH_KEY_UNNEST] = FileColumnMeaningType.UNNEST_KEY
         if eval_set is not None and len(eval_set) > 0:
             meaning_types[EVAL_SET_INDEX] = FileColumnMeaningType.EVAL_SET_INDEX
 
-        df = self.__add_fit_system_record_id(df, meaning_types, self.fit_search_keys)
+        df = self.__add_fit_system_record_id(df, self.fit_search_keys, SYSTEM_RECORD_ID)
+
+        if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
+            df = df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL)
+
+        meaning_types[SYSTEM_RECORD_ID] = FileColumnMeaningType.SYSTEM_RECORD_ID
 
-        self.df_with_original_index = df.copy()
         df = df.reset_index(drop=True).sort_values(by=SYSTEM_RECORD_ID).reset_index(drop=True)
 
         combined_search_keys = combine_search_keys(self.fit_search_keys.keys())
@@ -2353,16 +2699,16 @@ class FeaturesEnricher(TransformerMixin):
         dataset = Dataset(
             "tds_" + str(uuid.uuid4()),
             df=df,
-            model_task_type=model_task_type,
+            meaning_types=meaning_types,
+            search_keys=combined_search_keys,
+            unnest_search_keys=unnest_search_keys,
+            model_task_type=self.model_task_type,
             date_format=self.date_format,
             random_state=self.random_state,
             rest_client=self.rest_client,
             logger=self.logger,
         )
-        dataset.meaning_types = meaning_types
-        dataset.search_keys = combined_search_keys
-        if email_converted_to_hem:
-            dataset.ignore_columns = [email_column]
+        dataset.columns_renaming = self.fit_columns_renaming
 
         self.passed_features = [
             column for column, meaning_type in meaning_types.items() if meaning_type == FileColumnMeaningType.FEATURE
@@ -2438,9 +2784,7 @@ class FeaturesEnricher(TransformerMixin):
             zero_hit_columns = self.get_columns_by_search_keys(zero_hit_search_keys)
             if zero_hit_columns:
                 msg = self.bundle.get("features_info_zero_hit_rate_search_keys").format(zero_hit_columns)
-                self.logger.warning(msg)
-                self.__display_support_link(msg)
-                self.warning_counter.increment()
+                self.__log_warning(msg, show_support_link=True)
 
         if (
             self._search_task.unused_features_for_generation is not None
@@ -2450,9 +2794,7 @@ class FeaturesEnricher(TransformerMixin):
                 dataset.columns_renaming.get(col) or col for col in self._search_task.unused_features_for_generation
             ]
             msg = self.bundle.get("features_not_generated").format(unused_features_for_generation)
-            self.logger.warning(msg)
-            print(msg)
-            self.warning_counter.increment()
+            self.__log_warning(msg)
 
         self.__prepare_feature_importances(trace_id, validated_X.columns.to_list() + self.fit_generated_features)
 
@@ -2460,7 +2802,13 @@ class FeaturesEnricher(TransformerMixin):
 
         autofe_description = self.get_autofe_features_description()
         if autofe_description is not None:
-            display_html_dataframe(autofe_description, autofe_description, "*Description of AutoFE feature names")
+            self.logger.info(f"AutoFE descriptions: {autofe_description}")
+            self.autofe_features_display_handle = display_html_dataframe(
+                df=autofe_description,
+                internal_df=autofe_description,
+                header=self.bundle.get("autofe_descriptions_header"),
+                display_id="autofe_descriptions",
+            )
 
         if self._has_paid_features(exclude_features_sources):
             if calculate_metrics is not None and calculate_metrics:
@@ -2500,32 +2848,32 @@ class FeaturesEnricher(TransformerMixin):
                         progress_callback,
                     )
                 except Exception:
-                    self.__show_report_button()
+                    self.report_button_handle = self.__show_report_button(display_id="report_button")
                     raise
 
-        self.__show_report_button()
+        self.report_button_handle = self.__show_report_button(display_id="report_button")
 
         if not self.warning_counter.has_warnings():
             self.__display_support_link(self.bundle.get("all_ok_community_invite"))
 
-    def __adjust_cv(self, df: pd.DataFrame, date_column: pd.Series, model_task_type: ModelTaskType):
+    def __adjust_cv(self, df: pd.DataFrame):
+        date_column = SearchKey.find_key(self.fit_search_keys, [SearchKey.DATE, SearchKey.DATETIME])
         # Check Multivariate time series
         if (
             self.cv is None
             and date_column
-            and model_task_type == ModelTaskType.REGRESSION
+            and self.model_task_type == ModelTaskType.REGRESSION
             and len({SearchKey.PHONE, SearchKey.EMAIL, SearchKey.HEM}.intersection(self.fit_search_keys.keys())) == 0
             and is_blocked_time_series(df, date_column, list(self.fit_search_keys.keys()) + [TARGET])
         ):
             msg = self.bundle.get("multivariate_timeseries_detected")
             self.__override_cv(CVType.blocked_time_series, msg, print_warning=False)
-        elif (
-            self.cv is None
-            and model_task_type != ModelTaskType.REGRESSION
-            and self._get_group_columns(df, self.fit_search_keys)
-        ):
+        elif self.cv is None and self.model_task_type != ModelTaskType.REGRESSION:
             msg = self.bundle.get("group_k_fold_in_classification")
             self.__override_cv(CVType.group_k_fold, msg, print_warning=self.cv is not None)
+            group_columns = self._get_group_columns(df, self.fit_search_keys)
+            self.runtime_parameters.properties["cv_params.group_columns"] = ",".join(group_columns)
+            self.runtime_parameters.properties["cv_params.shuffle_kfold"] = "True"
 
     def __override_cv(self, cv: CVType, msg: str, print_warning: bool = True):
         if print_warning:
@@ -2543,9 +2891,6 @@ class FeaturesEnricher(TransformerMixin):
         return [c for c, v in search_keys_with_autodetection.items() if v.value.value in keys]
 
     def _validate_X(self, X, is_transform=False) -> pd.DataFrame:
-        if _num_samples(X) == 0:
-            raise ValidationError(self.bundle.get("x_is_empty"))
-
         if isinstance(X, pd.DataFrame):
             if isinstance(X.columns, pd.MultiIndex) or isinstance(X.index, pd.MultiIndex):
                 raise ValidationError(self.bundle.get("x_multiindex_unsupported"))
@@ -2559,6 +2904,9 @@ class FeaturesEnricher(TransformerMixin):
         else:
             raise ValidationError(self.bundle.get("unsupported_x_type").format(type(X)))
 
+        if _num_samples(X) == 0:
+            raise ValidationError(self.bundle.get("x_is_empty"))
+
         if len(set(validated_X.columns)) != len(validated_X.columns):
             raise ValidationError(self.bundle.get("x_contains_dup_columns"))
         if not is_transform and not validated_X.index.is_unique:
@@ -2578,13 +2926,12 @@ class FeaturesEnricher(TransformerMixin):
             raise ValidationError(self.bundle.get("x_contains_reserved_column_name").format(EVAL_SET_INDEX))
         if SYSTEM_RECORD_ID in validated_X.columns:
             raise ValidationError(self.bundle.get("x_contains_reserved_column_name").format(SYSTEM_RECORD_ID))
+        if ENTITY_SYSTEM_RECORD_ID in validated_X.columns:
+            raise ValidationError(self.bundle.get("x_contains_reserved_column_name").format(ENTITY_SYSTEM_RECORD_ID))
 
         return validated_X
 
     def _validate_y(self, X: pd.DataFrame, y) -> pd.Series:
-        if _num_samples(y) == 0:
-            raise ValidationError(self.bundle.get("y_is_empty"))
-
         if (
             not isinstance(y, pd.Series)
             and not isinstance(y, pd.DataFrame)
@@ -2593,6 +2940,9 @@ class FeaturesEnricher(TransformerMixin):
         ):
             raise ValidationError(self.bundle.get("unsupported_y_type").format(type(y)))
 
+        if _num_samples(y) == 0:
+            raise ValidationError(self.bundle.get("y_is_empty"))
+
         if _num_samples(X) != _num_samples(y):
             raise ValidationError(self.bundle.get("x_and_y_diff_size").format(_num_samples(X), _num_samples(y)))
 
@@ -2730,9 +3080,10 @@ class FeaturesEnricher(TransformerMixin):
         X: pd.DataFrame, y: pd.Series, cv: Optional[CVType]
     ) -> Tuple[pd.DataFrame, pd.Series]:
         if cv not in [CVType.time_series, CVType.blocked_time_series]:
+            record_id_column = ENTITY_SYSTEM_RECORD_ID if ENTITY_SYSTEM_RECORD_ID in X else SYSTEM_RECORD_ID
             Xy = X.copy()
             Xy[TARGET] = y
-            Xy = Xy.sort_values(by=SYSTEM_RECORD_ID).reset_index(drop=True)
+            Xy = Xy.sort_values(by=record_id_column).reset_index(drop=True)
             X = Xy.drop(columns=TARGET)
             y = Xy[TARGET].copy()
 
@@ -2750,7 +3101,7 @@ class FeaturesEnricher(TransformerMixin):
             if DateTimeSearchKeyConverter.DATETIME_COL in X.columns:
                 date_column = DateTimeSearchKeyConverter.DATETIME_COL
             else:
-                date_column = FeaturesEnricher._get_date_column(search_keys)
+                date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
             sort_columns = [date_column] if date_column is not None else []
 
             # Xy = pd.concat([X, y], axis=1)
@@ -2846,7 +3197,7 @@ class FeaturesEnricher(TransformerMixin):
 
             do_without_pandas_limits(print_datasets_sample)
 
-            maybe_date_col = self._get_date_column(self.search_keys)
+            maybe_date_col = SearchKey.find_key(self.search_keys, [SearchKey.DATE, SearchKey.DATETIME])
             if X is not None and maybe_date_col is not None and maybe_date_col in X.columns:
                 # TODO cast date column to single dtype
                 date_converter = DateTimeSearchKeyConverter(maybe_date_col, self.date_format)
@@ -2856,7 +3207,7 @@ class FeaturesEnricher(TransformerMixin):
                 self.logger.info(f"Dates interval is ({min_date}, {max_date})")
 
         except Exception:
-            self.logger.exception("Failed to log debug information")
+            self.logger.warning("Failed to log debug information", exc_info=True)
 
     def __handle_index_search_keys(self, df: pd.DataFrame, search_keys: Dict[str, SearchKey]) -> pd.DataFrame:
         index_names = df.index.names if df.index.names != [None] else [DEFAULT_INDEX]
@@ -2876,15 +3227,8 @@ class FeaturesEnricher(TransformerMixin):
 
         return df
 
-    @staticmethod
-    def _get_date_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
-        for col, t in search_keys.items():
-            if t in [SearchKey.DATE, SearchKey.DATETIME]:
-                return col
-
-    @staticmethod
     def _add_current_date_as_key(
-        df: pd.DataFrame, search_keys: Dict[str, SearchKey], logger: logging.Logger, bundle: ResourceBundle
+        self, df: pd.DataFrame, search_keys: Dict[str, SearchKey], logger: logging.Logger, bundle: ResourceBundle
     ) -> pd.DataFrame:
         if (
             set(search_keys.values()) == {SearchKey.PHONE}
@@ -2892,12 +3236,10 @@ class FeaturesEnricher(TransformerMixin):
             or set(search_keys.values()) == {SearchKey.HEM}
             or set(search_keys.values()) == {SearchKey.COUNTRY, SearchKey.POSTAL_CODE}
         ):
-            msg = bundle.get("current_date_added")
-            print(msg)
-            logger.warning(msg)
+            self.__log_warning(bundle.get("current_date_added"))
             df[FeaturesEnricher.CURRENT_DATE] = datetime.date.today()
             search_keys[FeaturesEnricher.CURRENT_DATE] = SearchKey.DATE
-            converter = DateTimeSearchKeyConverter(FeaturesEnricher.CURRENT_DATE, None, logger, bundle)
+            converter = DateTimeSearchKeyConverter(FeaturesEnricher.CURRENT_DATE)
             df = converter.convert(df)
         return df
 
@@ -2911,24 +3253,87 @@ class FeaturesEnricher(TransformerMixin):
 
     @staticmethod
     def _get_email_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+        cols = [col for col, t in search_keys.items() if t == SearchKey.EMAIL]
+        if len(cols) > 1:
+            raise Exception("More than one email column found after unnest")
+        if len(cols) == 1:
+            return cols[0]
+
+    @staticmethod
+    def _get_hem_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+        cols = [col for col, t in search_keys.items() if t == SearchKey.HEM]
+        if len(cols) > 1:
+            raise Exception("More than one hem column found after unnest")
+        if len(cols) == 1:
+            return cols[0]
+
+    @staticmethod
+    def _get_ip_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+        cols = [col for col, t in search_keys.items() if t == SearchKey.IP]
+        if len(cols) > 1:
+            raise Exception("More than one ip column found after unnest")
+        if len(cols) == 1:
+            return cols[0]
+
+    @staticmethod
+    def _get_phone_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
         for col, t in search_keys.items():
-            if t == SearchKey.EMAIL:
+            if t == SearchKey.PHONE:
                 return col
 
     @staticmethod
-    def _get_hem_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_country_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
         for col, t in search_keys.items():
-            if t == SearchKey.HEM:
+            if t == SearchKey.COUNTRY:
                 return col
 
     @staticmethod
-    def _get_phone_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_postal_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
         for col, t in search_keys.items():
-            if t == SearchKey.PHONE:
+            if t == SearchKey.POSTAL_CODE:
                 return col
 
+    def _explode_multiple_search_keys(
+        self, df: pd.DataFrame, search_keys: Dict[str, SearchKey], columns_renaming: Dict[str, str]
+    ) -> Tuple[pd.DataFrame, Dict[str, List[str]]]:
+        # find groups of multiple search keys
+        search_key_names_by_type: Dict[SearchKey, List[str]] = {}
+        for key_name, key_type in search_keys.items():
+            search_key_names_by_type[key_type] = search_key_names_by_type.get(key_type, []) + [key_name]
+        search_key_names_by_type = {
+            key_type: key_names for key_type, key_names in search_key_names_by_type.items() if len(key_names) > 1
+        }
+        if len(search_key_names_by_type) == 0:
+            return df, {}
+
+        self.logger.info(f"Start exploding dataset by {search_key_names_by_type}. Size before: {len(df)}")
+        multiple_keys_columns = [col for cols in search_key_names_by_type.values() for col in cols]
+        other_columns = [col for col in df.columns if col not in multiple_keys_columns]
+        exploded_dfs = []
+        unnest_search_keys = {}
+
+        for key_type, key_names in search_key_names_by_type.items():
+            new_search_key = f"upgini_{key_type.name.lower()}_unnest"
+            exploded_df = pd.melt(
+                df, id_vars=other_columns, value_vars=key_names, var_name=SEARCH_KEY_UNNEST, value_name=new_search_key
+            )
+            exploded_dfs.append(exploded_df)
+            for old_key in key_names:
+                del search_keys[old_key]
+            search_keys[new_search_key] = key_type
+            unnest_search_keys[new_search_key] = key_names
+            columns_renaming[new_search_key] = new_search_key
+
+        df = pd.concat(exploded_dfs, ignore_index=True)
+        self.logger.info(f"Finished explosion. Size after: {len(df)}")
+        return df, unnest_search_keys
+
     def __add_fit_system_record_id(
-        self, df: pd.DataFrame, meaning_types: Dict[str, FileColumnMeaningType], search_keys: Dict[str, SearchKey]
+        self,
+        df: pd.DataFrame,
+        # meaning_types: Dict[str, FileColumnMeaningType],
+        search_keys: Dict[str, SearchKey],
+        id_name: str,
     ) -> pd.DataFrame:
         # save original order or rows
         original_index_name = df.index.name
@@ -2939,52 +3344,61 @@ class FeaturesEnricher(TransformerMixin):
 
         # order by date and idempotent order by other keys
         if self.cv not in [CVType.time_series, CVType.blocked_time_series]:
-            sort_exclude_columns = [original_order_name, ORIGINAL_INDEX, EVAL_SET_INDEX, TARGET, "__target"]
+            sort_exclude_columns = [
+                original_order_name,
+                ORIGINAL_INDEX,
+                EVAL_SET_INDEX,
+                TARGET,
+                "__target",
+                ENTITY_SYSTEM_RECORD_ID,
+            ]
             if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
                 date_column = DateTimeSearchKeyConverter.DATETIME_COL
-                sort_exclude_columns.append(self._get_date_column(search_keys))
+                sort_exclude_columns.append(SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME]))
             else:
-                date_column = self._get_date_column(search_keys)
+                date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
             sort_columns = [date_column] if date_column is not None else []
 
+            sorted_other_keys = sorted(search_keys, key=lambda x: str(search_keys.get(x)))
+            sorted_other_keys = [k for k in sorted_other_keys if k not in sort_exclude_columns]
+
             other_columns = sorted(
                 [
                     c
                     for c in df.columns
-                    if c not in sort_columns and c not in sort_exclude_columns and df[c].nunique() > 1
+                    if c not in sort_columns
+                    and c not in sorted_other_keys
+                    and c not in sort_exclude_columns
+                    and df[c].nunique() > 1
                 ]
-                # [
-                #     sk
-                #     for sk, key_type in search_keys.items()
-                #     if key_type not in [SearchKey.DATE, SearchKey.DATETIME]
-                #     and sk in df.columns
-                #     and df[sk].nunique() > 1  # don't use constant keys for hash
-                # ]
             )
 
+            all_other_columns = sorted_other_keys + other_columns
+
             search_keys_hash = "search_keys_hash"
-            if len(other_columns) > 0:
+            if len(all_other_columns) > 0:
                 sort_columns.append(search_keys_hash)
-                df[search_keys_hash] = pd.util.hash_pandas_object(df[other_columns], index=False)
+                df[search_keys_hash] = pd.util.hash_pandas_object(df[all_other_columns], index=False)
 
             df = df.sort_values(by=sort_columns)
 
             if search_keys_hash in df.columns:
                 df.drop(columns=search_keys_hash, inplace=True)
 
-        if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
-            df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL, inplace=True)
-
         df = df.reset_index(drop=True).reset_index()
         # system_record_id saves correct order for fit
-        df = df.rename(columns={DEFAULT_INDEX: SYSTEM_RECORD_ID})
+        df = df.rename(columns={DEFAULT_INDEX: id_name})
 
         # return original order
         df = df.set_index(ORIGINAL_INDEX)
         df.index.name = original_index_name
         df = df.sort_values(by=original_order_name).drop(columns=original_order_name)
 
-        meaning_types[SYSTEM_RECORD_ID] = FileColumnMeaningType.SYSTEM_RECORD_ID
+        # meaning_types[id_name] = (
+        #     FileColumnMeaningType.SYSTEM_RECORD_ID
+        #     if id_name == SYSTEM_RECORD_ID
+        #     else FileColumnMeaningType.ENTITY_SYSTEM_RECORD_ID
+        # )
         return df
 
     def __correct_target(self, df: pd.DataFrame) -> pd.DataFrame:
@@ -3039,7 +3453,11 @@ class FeaturesEnricher(TransformerMixin):
         )
 
         comparing_columns = X.columns if is_transform else df_with_original_index.columns
-        dup_features = [c for c in comparing_columns if c in result_features.columns and c != SYSTEM_RECORD_ID]
+        dup_features = [
+            c
+            for c in comparing_columns
+            if c in result_features.columns and c not in [SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID]
+        ]
         if len(dup_features) > 0:
             self.logger.warning(f"X contain columns with same name as returned from backend: {dup_features}")
             raise ValidationError(self.bundle.get("returned_features_same_as_passed").format(dup_features))
@@ -3047,11 +3465,11 @@ class FeaturesEnricher(TransformerMixin):
         # index overrites from result_features
         original_index_name = df_with_original_index.index.name
         df_with_original_index = df_with_original_index.reset_index()
+        # TODO drop system_record_id before merge
         result_features = pd.merge(
             df_with_original_index,
             result_features,
-            left_on=SYSTEM_RECORD_ID,
-            right_on=SYSTEM_RECORD_ID,
+            on=ENTITY_SYSTEM_RECORD_ID,
             how="left" if is_transform else "inner",
         )
         result_features = result_features.set_index(original_index_name or DEFAULT_INDEX)
@@ -3059,10 +3477,12 @@ class FeaturesEnricher(TransformerMixin):
 
         if rows_to_drop is not None:
             self.logger.info(f"Before dropping target outliers size: {len(result_features)}")
-            result_features = result_features[~result_features[SYSTEM_RECORD_ID].isin(rows_to_drop[SYSTEM_RECORD_ID])]
+            result_features = result_features[
+                ~result_features[ENTITY_SYSTEM_RECORD_ID].isin(rows_to_drop[ENTITY_SYSTEM_RECORD_ID])
+            ]
             self.logger.info(f"After dropping target outliers size: {len(result_features)}")
 
-        result_eval_sets = dict()
+        result_eval_sets = {}
         if not is_transform and EVAL_SET_INDEX in result_features.columns:
             result_train_features = result_features.loc[result_features[EVAL_SET_INDEX] == 0].copy()
             eval_set_indices = list(result_features[EVAL_SET_INDEX].unique())
@@ -3092,16 +3512,17 @@ class FeaturesEnricher(TransformerMixin):
             result_train = result_train_features
 
         if drop_system_record_id:
-            if SYSTEM_RECORD_ID in result_train.columns:
-                result_train = result_train.drop(columns=SYSTEM_RECORD_ID)
+            result_train = result_train.drop(columns=[SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID], errors="ignore")
             for eval_set_index in result_eval_sets.keys():
-                if SYSTEM_RECORD_ID in result_eval_sets[eval_set_index].columns:
-                    result_eval_sets[eval_set_index] = result_eval_sets[eval_set_index].drop(columns=SYSTEM_RECORD_ID)
+                result_eval_sets[eval_set_index] = result_eval_sets[eval_set_index].drop(
+                    columns=[SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID], errors="ignore"
+                )
 
         return result_train, result_eval_sets
 
-    def __prepare_feature_importances(self, trace_id: str, x_columns: List[str], silent=False):
-        llm_source = "LLM with external data augmentation"
+    def __prepare_feature_importances(
+            self, trace_id: str, x_columns: List[str], updated_shaps: Optional[Dict[str, float]] = None, silent=False
+    ):
         if self._search_task is None:
             raise NotFittedError(self.bundle.get("transform_unfitted_enricher"))
         features_meta = self._search_task.get_all_features_metadata_v2()
@@ -3112,122 +3533,44 @@ class FeaturesEnricher(TransformerMixin):
         features_df = self._search_task.get_all_initial_raw_features(trace_id, metrics_calculation=True)
 
         self.feature_names_ = []
+        self.dropped_client_feature_names_ = []
         self.feature_importances_ = []
         features_info = []
         features_info_without_links = []
         internal_features_info = []
 
-        def round_shap_value(shap: float) -> float:
-            if shap > 0.0 and shap < 0.0001:
-                return 0.0001
-            else:
-                return round(shap, 4)
-
-        def list_or_single(lst: List[str], single: str):
-            return lst or ([single] if single else [])
-
-        def to_anchor(link: str, value: str) -> str:
-            if not value:
-                return ""
-            elif not link:
-                return value
-            elif value == llm_source:
-                return value
-            else:
-                return f"<a href='{link}' target='_blank' rel='noopener noreferrer'>{value}</a>"
-
-        def make_links(names: List[str], links: List[str]):
-            all_links = [to_anchor(link, name) for name, link in itertools.zip_longest(names, links)]
-            return ",".join(all_links)
+        if updated_shaps is not None:
+            for fm in features_meta:
+                fm.shap_value = updated_shaps.get(fm.name, 0.0)
 
         features_meta.sort(key=lambda m: (-m.shap_value, m.name))
         for feature_meta in features_meta:
             if feature_meta.name in original_names_dict.keys():
                 feature_meta.name = original_names_dict[feature_meta.name]
-            # Use only enriched features
+
+            is_client_feature = feature_meta.name in x_columns
+
+            if feature_meta.shap_value == 0.0:
+                if self.select_features:
+                    self.dropped_client_feature_names_.append(feature_meta.name)
+                continue
+
+            # Use only important features
             if (
-                feature_meta.name in x_columns
+                feature_meta.name in self.fit_generated_features
                 or feature_meta.name == COUNTRY
-                or feature_meta.shap_value == 0.0
-                or feature_meta.name in self.fit_generated_features
+                # In select_features mode we select also from etalon features and need to show them
+                or (not self.select_features and is_client_feature)
             ):
                 continue
 
-            feature_sample = []
             self.feature_names_.append(feature_meta.name)
-            self.feature_importances_.append(round_shap_value(feature_meta.shap_value))
-            if feature_meta.name in features_df.columns:
-                feature_sample = np.random.choice(features_df[feature_meta.name].dropna().unique(), 3).tolist()
-                if len(feature_sample) > 0 and isinstance(feature_sample[0], float):
-                    feature_sample = [round(f, 4) for f in feature_sample]
-                feature_sample = [str(f) for f in feature_sample]
-                feature_sample = ", ".join(feature_sample)
-                if len(feature_sample) > 30:
-                    feature_sample = feature_sample[:30] + "..."
-
-            internal_provider = feature_meta.data_provider or "Upgini"
-            providers = list_or_single(feature_meta.data_providers, feature_meta.data_provider)
-            provider_links = list_or_single(feature_meta.data_provider_links, feature_meta.data_provider_link)
-            if providers:
-                provider = make_links(providers, provider_links)
-            else:
-                provider = to_anchor("https://upgini.com", "Upgini")
+            self.feature_importances_.append(_round_shap_value(feature_meta.shap_value))
 
-            internal_source = feature_meta.data_source or (
-                llm_source
-                if not feature_meta.name.endswith("_country") and not feature_meta.name.endswith("_postal_code")
-                else ""
-            )
-            sources = list_or_single(feature_meta.data_sources, feature_meta.data_source)
-            source_links = list_or_single(feature_meta.data_source_links, feature_meta.data_source_link)
-            if sources:
-                source = make_links(sources, source_links)
-            else:
-                source = internal_source
-
-            internal_feature_name = feature_meta.name
-            if feature_meta.doc_link:
-                feature_name = to_anchor(feature_meta.doc_link, feature_meta.name)
-            else:
-                feature_name = internal_feature_name
-
-            features_info.append(
-                {
-                    self.bundle.get("features_info_name"): feature_name,
-                    self.bundle.get("features_info_shap"): round_shap_value(feature_meta.shap_value),
-                    self.bundle.get("features_info_hitrate"): feature_meta.hit_rate,
-                    self.bundle.get("features_info_value_preview"): feature_sample,
-                    self.bundle.get("features_info_provider"): provider,
-                    self.bundle.get("features_info_source"): source,
-                    self.bundle.get("features_info_update_frequency"): feature_meta.update_frequency,
-                }
-            )
-            features_info_without_links.append(
-                {
-                    self.bundle.get("features_info_name"): internal_feature_name,
-                    self.bundle.get("features_info_shap"): round_shap_value(feature_meta.shap_value),
-                    self.bundle.get("features_info_hitrate"): feature_meta.hit_rate,
-                    self.bundle.get("features_info_value_preview"): feature_sample,
-                    self.bundle.get("features_info_provider"): internal_provider,
-                    self.bundle.get("features_info_source"): internal_source,
-                    self.bundle.get("features_info_update_frequency"): feature_meta.update_frequency,
-                }
-            )
-            internal_features_info.append(
-                {
-                    self.bundle.get("features_info_name"): internal_feature_name,
-                    "feature_link": feature_meta.doc_link,
-                    self.bundle.get("features_info_shap"): round_shap_value(feature_meta.shap_value),
-                    self.bundle.get("features_info_hitrate"): feature_meta.hit_rate,
-                    self.bundle.get("features_info_value_preview"): feature_sample,
-                    self.bundle.get("features_info_provider"): internal_provider,
-                    "provider_link": feature_meta.data_provider_link,
-                    self.bundle.get("features_info_source"): internal_source,
-                    "source_link": feature_meta.data_source_link,
-                    self.bundle.get("features_info_commercial_schema"): feature_meta.commercial_schema or "",
-                    self.bundle.get("features_info_update_frequency"): feature_meta.update_frequency,
-                }
-            )
+            feature_info = FeatureInfo.from_metadata(feature_meta, features_df, is_client_feature)
+            features_info.append(feature_info.to_row(self.bundle))
+            features_info_without_links.append(feature_info.to_row_without_links(self.bundle))
+            internal_features_info.append(feature_info.to_internal_row(self.bundle))
 
         if len(features_info) > 0:
             self.features_info = pd.DataFrame(features_info)
@@ -3252,7 +3595,22 @@ class FeaturesEnricher(TransformerMixin):
             autofe_meta = self._search_task.get_autofe_metadata()
             if autofe_meta is None:
                 return None
-            features_meta = self._search_task.get_all_features_metadata_v2()
+            if len(self._internal_features_info) != 0:
+
+                def to_feature_meta(row):
+                    fm = FeaturesMetadataV2(
+                        name=row[bundle.get("features_info_name")],
+                        type="",
+                        source="",
+                        hit_rate=row[bundle.get("features_info_hitrate")],
+                        shap_value=row[bundle.get("features_info_shap")],
+                        data_source=row[bundle.get("features_info_source")],
+                    )
+                    return fm
+
+                features_meta = self._internal_features_info.apply(to_feature_meta, axis=1).to_list()
+            else:
+                features_meta = self._search_task.get_all_features_metadata_v2()
 
             def get_feature_by_name(name: str):
                 for m in features_meta:
@@ -3261,41 +3619,52 @@ class FeaturesEnricher(TransformerMixin):
 
             descriptions = []
             for m in autofe_meta:
-                autofe_feature = Feature.from_formula(m.formula)
                 orig_to_hashed = {base_column.original_name: base_column.hashed_name for base_column in m.base_columns}
-                autofe_feature.rename_columns(orig_to_hashed)
-                autofe_feature.set_display_index(m.display_index)
+
+                autofe_feature = (
+                    Feature.from_formula(m.formula)
+                    .set_display_index(m.display_index)
+                    .set_alias(m.alias)
+                    .set_op_params(m.operator_params or {})
+                    .rename_columns(orig_to_hashed)
+                )
+
                 if autofe_feature.op.is_vector:
                     continue
 
-                description = dict()
+                description = {}
 
                 feature_meta = get_feature_by_name(autofe_feature.get_display_name(shorten=True))
                 if feature_meta is None:
                     self.logger.warning(f"Feature meta for display index {m.display_index} not found")
                     continue
                 description["shap"] = feature_meta.shap_value
-                description["Sources"] = feature_meta.data_source.replace("AutoFE: features from ", "").replace(
-                    "AutoFE: feature from ", ""
-                )
-                description["Feature name"] = feature_meta.name
+                description[self.bundle.get("autofe_descriptions_sources")] = feature_meta.data_source.replace(
+                    "AutoFE: features from ", ""
+                ).replace("AutoFE: feature from ", "")
+                description[self.bundle.get("autofe_descriptions_feature_name")] = feature_meta.name
 
                 feature_idx = 1
                 for bc in m.base_columns:
-                    description[f"Feature {feature_idx}"] = bc.hashed_name
+                    description[self.bundle.get("autofe_descriptions_feature").format(feature_idx)] = bc.hashed_name
                     feature_idx += 1
 
-                description["Function"] = autofe_feature.op.name
+                description[self.bundle.get("autofe_descriptions_function")] = ",".join(
+                    sorted(autofe_feature.get_all_operand_names())
+                )
 
                 descriptions.append(description)
 
             if len(descriptions) == 0:
                 return None
 
-            descriptions_df = pd.DataFrame(descriptions)
-            descriptions_df.fillna("", inplace=True)
-            descriptions_df.sort_values(by="shap", ascending=False, inplace=True)
-            descriptions_df.drop(columns="shap", inplace=True)
+            descriptions_df = (
+                pd.DataFrame(descriptions)
+                .fillna("")
+                .sort_values(by="shap", ascending=False)
+                .drop(columns="shap")
+                .reset_index(drop=True)
+            )
             return descriptions_df
 
         except Exception:
@@ -3348,10 +3717,16 @@ class FeaturesEnricher(TransformerMixin):
         is_transform=False,
         silent_mode=False,
     ):
+        for _, key_type in search_keys.items():
+            if not isinstance(key_type, SearchKey):
+                raise ValidationError(self.bundle.get("unsupported_type_of_search_key").format(key_type))
+
         valid_search_keys = {}
         unsupported_search_keys = {
             SearchKey.IP_RANGE_FROM,
             SearchKey.IP_RANGE_TO,
+            SearchKey.IPV6_RANGE_FROM,
+            SearchKey.IPV6_RANGE_TO,
             SearchKey.MSISDN_RANGE_FROM,
             SearchKey.MSISDN_RANGE_TO,
             # SearchKey.EMAIL_ONE_DOMAIN,
@@ -3360,11 +3735,17 @@ class FeaturesEnricher(TransformerMixin):
         if len(passed_unsupported_search_keys) > 0:
             raise ValidationError(self.bundle.get("unsupported_search_key").format(passed_unsupported_search_keys))
 
+        x_columns = [
+            c
+            for c in x.columns
+            if c not in [TARGET, EVAL_SET_INDEX, SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID, SEARCH_KEY_UNNEST]
+        ]
+
         for column_id, meaning_type in search_keys.items():
             column_name = None
             if isinstance(column_id, str):
                 if column_id not in x.columns:
-                    raise ValidationError(self.bundle.get("search_key_not_found").format(column_id, list(x.columns)))
+                    raise ValidationError(self.bundle.get("search_key_not_found").format(column_id, x_columns))
                 column_name = column_id
                 valid_search_keys[column_name] = meaning_type
             elif isinstance(column_id, int):
@@ -3378,15 +3759,15 @@ class FeaturesEnricher(TransformerMixin):
             if meaning_type == SearchKey.COUNTRY and self.country_code is not None:
                 msg = self.bundle.get("search_key_country_and_country_code")
                 self.logger.warning(msg)
-                print(msg)
+                if not silent_mode:
+                    self.__log_warning(msg)
                 self.country_code = None
 
             if not self.__is_registered and not is_demo_dataset and meaning_type in SearchKey.personal_keys():
                 msg = self.bundle.get("unregistered_with_personal_keys").format(meaning_type)
                 self.logger.warning(msg)
                 if not silent_mode:
-                    self.warning_counter.increment()
-                    print(msg)
+                    self.__log_warning(msg)
 
                 valid_search_keys[column_name] = SearchKey.CUSTOM_KEY
             else:
@@ -3420,27 +3801,23 @@ class FeaturesEnricher(TransformerMixin):
             and not silent_mode
         ):
             msg = self.bundle.get("date_only_search")
-            print(msg)
-            self.logger.warning(msg)
-            self.warning_counter.increment()
+            self.__log_warning(msg)
 
         maybe_date = [k for k, v in valid_search_keys.items() if v in [SearchKey.DATE, SearchKey.DATETIME]]
         if (self.cv is None or self.cv == CVType.k_fold) and len(maybe_date) > 0 and not silent_mode:
             date_column = next(iter(maybe_date))
             if x[date_column].nunique() > 0.9 * _num_samples(x):
                 msg = self.bundle.get("date_search_without_time_series")
-                print(msg)
-                self.logger.warning(msg)
-                self.warning_counter.increment()
+                self.__log_warning(msg)
 
         if len(valid_search_keys) == 1:
-            for k, v in valid_search_keys.items():
-                # Show warning for country only if country is the only key
-                if x[k].nunique() == 1 and (v != SearchKey.COUNTRY or len(valid_search_keys) == 1):
-                    msg = self.bundle.get("single_constant_search_key").format(v, x[k].values[0])
-                    print(msg)
-                    self.logger.warning(msg)
-                    self.warning_counter.increment()
+            key, value = list(valid_search_keys.items())[0]
+            # Show warning for country only if country is the only key
+            if x[key].nunique() == 1:
+                msg = self.bundle.get("single_constant_search_key").format(value, x[key].values[0])
+                if not silent_mode:
+                    self.__log_warning(msg)
+                # TODO maybe raise ValidationError
 
         self.logger.info(f"Prepared search keys: {valid_search_keys}")
 
@@ -3473,7 +3850,10 @@ class FeaturesEnricher(TransformerMixin):
             display_html_dataframe(self.metrics, self.metrics, msg)
 
     def __show_selected_features(self, search_keys: Dict[str, SearchKey]):
-        msg = self.bundle.get("features_info_header").format(len(self.feature_names_), list(search_keys.keys()))
+        search_key_names = search_keys.keys()
+        if self.fit_columns_renaming:
+            search_key_names = [self.fit_columns_renaming.get(col, col) for col in search_key_names]
+        msg = self.bundle.get("features_info_header").format(len(self.feature_names_), search_key_names)
 
         try:
             _ = get_ipython()  # type: ignore
@@ -3481,27 +3861,29 @@ class FeaturesEnricher(TransformerMixin):
             print(Format.GREEN + Format.BOLD + msg + Format.END)
             self.logger.info(msg)
             if len(self.feature_names_) > 0:
-                display_html_dataframe(
-                    self.features_info, self._features_info_without_links, self.bundle.get("relevant_features_header")
+                self.features_info_display_handle = display_html_dataframe(
+                    self.features_info,
+                    self._features_info_without_links,
+                    self.bundle.get("relevant_features_header"),
+                    display_id="features_info",
                 )
 
-                display_html_dataframe(
+                self.data_sources_display_handle = display_html_dataframe(
                     self.relevant_data_sources,
                     self._relevant_data_sources_wo_links,
                     self.bundle.get("relevant_data_sources_header"),
+                    display_id="data_sources",
                 )
             else:
                 msg = self.bundle.get("features_info_zero_important_features")
-                self.logger.warning(msg)
-                self.__display_support_link(msg)
-                self.warning_counter.increment()
+                self.__log_warning(msg, show_support_link=True)
         except (ImportError, NameError):
             print(msg)
             print(self._internal_features_info)
 
-    def __show_report_button(self):
+    def __show_report_button(self, display_id: Optional[str] = None, display_handle=None):
         try:
-            prepare_and_show_report(
+            return prepare_and_show_report(
                 relevant_features_df=self._features_info_without_links,
                 relevant_datasources_df=self.relevant_data_sources,
                 metrics_df=self.metrics,
@@ -3509,6 +3891,8 @@ class FeaturesEnricher(TransformerMixin):
                 search_id=self._search_task.search_task_id,
                 email=self.rest_client.get_current_email(),
                 search_keys=[str(sk) for sk in self.search_keys.values()],
+                display_id=display_id,
+                display_handle=display_handle,
             )
         except Exception:
             pass
@@ -3550,65 +3934,70 @@ class FeaturesEnricher(TransformerMixin):
         def check_need_detect(search_key: SearchKey):
             return not is_transform or search_key in self.fit_search_keys.values()
 
-        if SearchKey.POSTAL_CODE not in search_keys.values() and check_need_detect(SearchKey.POSTAL_CODE):
-            maybe_key = PostalCodeSearchKeyDetector().get_search_key_column(sample)
-            if maybe_key is not None:
-                search_keys[maybe_key] = SearchKey.POSTAL_CODE
-                self.autodetected_search_keys[maybe_key] = SearchKey.POSTAL_CODE
-                self.logger.info(f"Autodetected search key POSTAL_CODE in column {maybe_key}")
+        # if SearchKey.POSTAL_CODE not in search_keys.values() and check_need_detect(SearchKey.POSTAL_CODE):
+        if check_need_detect(SearchKey.POSTAL_CODE):
+            maybe_keys = PostalCodeSearchKeyDetector().get_search_key_columns(sample, search_keys)
+            if maybe_keys:
+                new_keys = {key: SearchKey.POSTAL_CODE for key in maybe_keys}
+                search_keys.update(new_keys)
+                self.autodetected_search_keys.update(new_keys)
+                self.logger.info(f"Autodetected search key POSTAL_CODE in column {maybe_keys}")
                 if not silent_mode:
-                    print(self.bundle.get("postal_code_detected").format(maybe_key))
+                    print(self.bundle.get("postal_code_detected").format(maybe_keys))
 
         if (
             SearchKey.COUNTRY not in search_keys.values()
             and self.country_code is None
             and check_need_detect(SearchKey.COUNTRY)
         ):
-            maybe_key = CountrySearchKeyDetector().get_search_key_column(sample)
-            if maybe_key is not None:
-                search_keys[maybe_key] = SearchKey.COUNTRY
-                self.autodetected_search_keys[maybe_key] = SearchKey.COUNTRY
+            maybe_key = CountrySearchKeyDetector().get_search_key_columns(sample, search_keys)
+            if maybe_key:
+                search_keys[maybe_key[0]] = SearchKey.COUNTRY
+                self.autodetected_search_keys[maybe_key[0]] = SearchKey.COUNTRY
                 self.logger.info(f"Autodetected search key COUNTRY in column {maybe_key}")
                 if not silent_mode:
                     print(self.bundle.get("country_detected").format(maybe_key))
 
         if (
-            SearchKey.EMAIL not in search_keys.values()
-            and SearchKey.HEM not in search_keys.values()
+            # SearchKey.EMAIL not in search_keys.values()
+            SearchKey.HEM not in search_keys.values()
             and check_need_detect(SearchKey.HEM)
         ):
-            maybe_key = EmailSearchKeyDetector().get_search_key_column(sample)
-            if maybe_key is not None and maybe_key not in search_keys.keys():
+            maybe_keys = EmailSearchKeyDetector().get_search_key_columns(sample, search_keys)
+            if maybe_keys:
                 if self.__is_registered or is_demo_dataset:
-                    search_keys[maybe_key] = SearchKey.EMAIL
-                    self.autodetected_search_keys[maybe_key] = SearchKey.EMAIL
-                    self.logger.info(f"Autodetected search key EMAIL in column {maybe_key}")
+                    new_keys = {key: SearchKey.EMAIL for key in maybe_keys}
+                    search_keys.update(new_keys)
+                    self.autodetected_search_keys.update(new_keys)
+                    self.logger.info(f"Autodetected search key EMAIL in column {maybe_keys}")
                     if not silent_mode:
-                        print(self.bundle.get("email_detected").format(maybe_key))
+                        print(self.bundle.get("email_detected").format(maybe_keys))
                 else:
                     self.logger.warning(
-                        f"Autodetected search key EMAIL in column {maybe_key}. But not used because not registered user"
+                        f"Autodetected search key EMAIL in column {maybe_keys}."
+                        " But not used because not registered user"
                     )
                     if not silent_mode:
-                        print(self.bundle.get("email_detected_not_registered").format(maybe_key))
-                    self.warning_counter.increment()
+                        self.__log_warning(self.bundle.get("email_detected_not_registered").format(maybe_keys))
 
-        if SearchKey.PHONE not in search_keys.values() and check_need_detect(SearchKey.PHONE):
-            maybe_key = PhoneSearchKeyDetector().get_search_key_column(sample)
-            if maybe_key is not None and maybe_key not in search_keys.keys():
+        # if SearchKey.PHONE not in search_keys.values() and check_need_detect(SearchKey.PHONE):
+        if check_need_detect(SearchKey.PHONE):
+            maybe_keys = PhoneSearchKeyDetector().get_search_key_columns(sample, search_keys)
+            if maybe_keys:
                 if self.__is_registered or is_demo_dataset:
-                    search_keys[maybe_key] = SearchKey.PHONE
-                    self.autodetected_search_keys[maybe_key] = SearchKey.PHONE
-                    self.logger.info(f"Autodetected search key PHONE in column {maybe_key}")
+                    new_keys = {key: SearchKey.PHONE for key in maybe_keys}
+                    search_keys.update(new_keys)
+                    self.autodetected_search_keys.update(new_keys)
+                    self.logger.info(f"Autodetected search key PHONE in column {maybe_keys}")
                     if not silent_mode:
-                        print(self.bundle.get("phone_detected").format(maybe_key))
+                        print(self.bundle.get("phone_detected").format(maybe_keys))
                 else:
                     self.logger.warning(
-                        f"Autodetected search key PHONE in column {maybe_key}. But not used because not registered user"
+                        f"Autodetected search key PHONE in column {maybe_keys}. "
+                        "But not used because not registered user"
                     )
                     if not silent_mode:
-                        print(self.bundle.get("phone_detected_not_registered"))
-                    self.warning_counter.increment()
+                        self.__log_warning(self.bundle.get("phone_detected_not_registered"))
 
         return search_keys
 
@@ -3630,21 +4019,19 @@ class FeaturesEnricher(TransformerMixin):
         half_train = round(len(train) / 2)
         part1 = train[:half_train]
         part2 = train[half_train:]
-        train_psi = calculate_psi(part1[self.TARGET_NAME], part2[self.TARGET_NAME])
-        if train_psi > 0.2:
-            self.warning_counter.increment()
-            msg = self.bundle.get("train_unstable_target").format(train_psi)
-            print(msg)
-            self.logger.warning(msg)
+        train_psi_result = calculate_psi(part1[self.TARGET_NAME], part2[self.TARGET_NAME])
+        if isinstance(train_psi_result, Exception):
+            self.logger.exception("Failed to calculate train PSI", train_psi_result)
+        elif train_psi_result > 0.2:
+            self.__log_warning(self.bundle.get("train_unstable_target").format(train_psi_result))
 
         # 2. Check train-test PSI
         if eval1 is not None:
-            train_test_psi = calculate_psi(train[self.TARGET_NAME], eval1[self.TARGET_NAME])
-            if train_test_psi > 0.2:
-                self.warning_counter.increment()
-                msg = self.bundle.get("eval_unstable_target").format(train_test_psi)
-                print(msg)
-                self.logger.warning(msg)
+            train_test_psi_result = calculate_psi(train[self.TARGET_NAME], eval1[self.TARGET_NAME])
+            if isinstance(train_test_psi_result, Exception):
+                self.logger.exception("Failed to calculate test PSI", train_test_psi_result)
+            elif train_test_psi_result > 0.2:
+                self.__log_warning(self.bundle.get("eval_unstable_target").format(train_test_psi_result))
 
     def _dump_python_libs(self):
         try:
@@ -3666,8 +4053,8 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(f"Showing support link: {link_text}")
             display(
                 HTML(
-                    f"""<br/>{link_text} <a href='{support_link}' target='_blank' rel='noopener noreferrer'>
-                    here</a>"""
+                    f"""{link_text} <a href='{support_link}' target='_blank' rel='noopener noreferrer'>
+                    here</a><br/>"""
                 )
             )
         except (ImportError, NameError):
@@ -3712,7 +4099,7 @@ class FeaturesEnricher(TransformerMixin):
                     if y is not None:
                         with open(f"{tmp_dir}/y.pickle", "wb") as y_file:
                             pickle.dump(sample(y, xy_sample_index), y_file)
-                        if eval_set:
+                        if eval_set and _num_samples(eval_set[0][0]) > 0:
                             eval_xy_sample_index = rnd.randint(0, _num_samples(eval_set[0][0]), size=1000)
                             with open(f"{tmp_dir}/eval_x.pickle", "wb") as eval_x_file:
                                 pickle.dump(sample(eval_set[0][0], eval_xy_sample_index), eval_x_file)
@@ -3803,6 +4190,8 @@ def hash_input(X: pd.DataFrame, y: Optional[pd.Series] = None, eval_set: Optiona
         if y is not None:
             hashed_objects.append(pd.util.hash_pandas_object(y, index=False).values)
         if eval_set is not None:
+            if isinstance(eval_set, tuple):
+                eval_set = [eval_set]
             for eval_X, eval_y in eval_set:
                 hashed_objects.append(pd.util.hash_pandas_object(eval_X, index=False).values)
                 hashed_objects.append(pd.util.hash_pandas_object(eval_y, index=False).values)