upgini 1.2.124__py3-none-any.whl → 1.2.125__py3-none-any.whl

upgini/__about__.py

@@ -1 +1 @@
-__version__ = "1.2.124"
+__version__ = "1.2.125"

upgini/autofe/binary.py

@@ -1,5 +1,6 @@
 import abc
 from typing import Optional
+
 import Levenshtein
 import numpy as np
 import pandas as pd
@@ -201,7 +202,7 @@ class JaroWinklerSim1(StringSim):
     has_symmetry_importance: bool = True
 
     def _prepare_value(self, value: Optional[str]) -> Optional[str]:
-        return value
+        return value if value is not None and len(value) > 0 else None
 
     def _similarity(self, left: str, right: str) -> float:
         return jarowinkler_similarity(left, right)
@@ -216,7 +217,7 @@ class JaroWinklerSim2(StringSim):
     has_symmetry_importance: bool = True
 
     def _prepare_value(self, value: Optional[str]) -> Optional[str]:
-        return value[::-1] if value is not None else None
+        return value[::-1] if value is not None and len(value) > 0 else None
 
     def _similarity(self, left: str, right: str) -> float:
         return jarowinkler_similarity(left, right)
@@ -231,7 +232,7 @@ class LevenshteinSim(StringSim):
     has_symmetry_importance: bool = True
 
     def _prepare_value(self, value: Optional[str]) -> Optional[str]:
-        return value
+        return value if value is not None and len(value) > 0 else None
 
     def _similarity(self, left: str, right: str) -> float:
         return 1 - Levenshtein.distance(left, right) / max(len(left), len(right))

upgini/dataset.py

@@ -151,7 +151,9 @@ class Dataset:
     def etalon_def_checked(self) -> Dict[str, str]:
         if self.etalon_def is None:
             self.etalon_def = {
-                v.value: k for k, v in self.meaning_types_checked.items() if v != FileColumnMeaningType.FEATURE
+                v.value: k
+                for k, v in self.meaning_types_checked.items()
+                if v not in [FileColumnMeaningType.FEATURE, FileColumnMeaningType.DATE_FEATURE]
             }
 
         return self.etalon_def

upgini/features_enricher.py

@@ -76,7 +76,7 @@ from upgini.utils.custom_loss_utils import (
 )
 from upgini.utils.cv_utils import CVConfig, get_groups
 from upgini.utils.datetime_utils import (
-    DateTimeSearchKeyConverter,
+    DateTimeConverter,
     is_blocked_time_series,
     is_dates_distribution_valid,
     is_time_series,
@@ -220,7 +220,9 @@ class FeaturesEnricher(TransformerMixin):
         cv: CVType | None = None,
         loss: str | None = None,
         autodetect_search_keys: bool = True,
+        # deprecated, use text_features instead
         generate_features: list[str] | None = None,
+        text_features: list[str] | None = None,
         columns_for_online_api: list[str] | None = None,
         round_embeddings: int | None = None,
         logs_enabled: bool = True,
@@ -305,10 +307,8 @@ class FeaturesEnricher(TransformerMixin):
             search_task = SearchTask(search_id, rest_client=self.rest_client, logger=self.logger)
 
             print(self.bundle.get("search_by_task_id_start"))
-            trace_id = str(uuid.uuid4())
-            if self.print_trace_id:
-                print(f"https://app.datadoghq.eu/logs?query=%40trace_id%3A{trace_id}")
-            with MDC(trace_id=trace_id):
+            trace_id = time.time_ns()
+            with MDC(correlation_id=trace_id):
                 try:
                     self.logger.debug(f"FeaturesEnricher created from existing search: {search_id}")
                     self._search_task = search_task.poll_result(trace_id, quiet=True, check_fit=True)
@@ -342,14 +342,14 @@ class FeaturesEnricher(TransformerMixin):
         self.shared_datasets = shared_datasets
         if shared_datasets is not None:
             self.runtime_parameters.properties["shared_datasets"] = ",".join(shared_datasets)
-        self.generate_features = generate_features
+        self.generate_features = text_features or generate_features
         self.round_embeddings = round_embeddings
-        if generate_features is not None:
-            if len(generate_features) > self.GENERATE_FEATURES_LIMIT:
+        if self.generate_features is not None:
+            if len(self.generate_features) > self.GENERATE_FEATURES_LIMIT:
                 msg = self.bundle.get("too_many_generate_features").format(self.GENERATE_FEATURES_LIMIT)
                 self.logger.error(msg)
                 raise ValidationError(msg)
-            self.runtime_parameters.properties["generate_features"] = ",".join(generate_features)
+            self.runtime_parameters.properties["generate_features"] = ",".join(self.generate_features)
             if round_embeddings is not None:
                 if not isinstance(round_embeddings, int) or round_embeddings < 0:
                     msg = self.bundle.get("invalid_round_embeddings")
@@ -484,9 +484,9 @@ class FeaturesEnricher(TransformerMixin):
         stability_agg_func: str, optional (default="max")
             Function to aggregate stability values. Can be "max", "min", "mean".
         """
-        trace_id = str(uuid.uuid4())
+        trace_id = time.time_ns()
         if self.print_trace_id:
-            print(f"https://app.datadoghq.eu/logs?query=%40trace_id%3A{trace_id}")
+            print(f"https://app.datadoghq.eu/logs?query=%40correlation_id%3A{trace_id}")
         start_time = time.time()
         auto_fe_parameters = AutoFEParameters() if auto_fe_parameters is None else auto_fe_parameters
         search_progress = SearchProgress(0.0, ProgressStage.START_FIT)
@@ -498,7 +498,7 @@ class FeaturesEnricher(TransformerMixin):
             progress_bar.progress = search_progress.to_progress_bar()
             progress_bar.display()
 
-        with MDC(trace_id=trace_id):
+        with MDC(correlation_id=trace_id):
             if len(args) > 0:
                 msg = f"WARNING: Unsupported positional arguments for fit: {args}"
                 self.logger.warning(msg)
@@ -643,11 +643,11 @@ class FeaturesEnricher(TransformerMixin):
 
         self.warning_counter.reset()
         auto_fe_parameters = AutoFEParameters() if auto_fe_parameters is None else auto_fe_parameters
-        trace_id = str(uuid.uuid4())
+        trace_id = time.time_ns()
         if self.print_trace_id:
-            print(f"https://app.datadoghq.eu/logs?query=%40trace_id%3A{trace_id}")
+            print(f"https://app.datadoghq.eu/logs?query=%40correlation_id%3A{trace_id}")
         start_time = time.time()
-        with MDC(trace_id=trace_id):
+        with MDC(correlation_id=trace_id):
             if len(args) > 0:
                 msg = f"WARNING: Unsupported positional arguments for fit_transform: {args}"
                 self.logger.warning(msg)
@@ -745,8 +745,8 @@ class FeaturesEnricher(TransformerMixin):
     def transform(
         self,
         X: pd.DataFrame,
-        *args,
         y: pd.Series | None = None,
+        *args,
         exclude_features_sources: list[str] | None = None,
         keep_input: bool = True,
         trace_id: str | None = None,
@@ -787,9 +787,11 @@ class FeaturesEnricher(TransformerMixin):
             progress_bar.progress = search_progress.to_progress_bar()
             if new_progress:
                 progress_bar.display()
-        trace_id = trace_id or str(uuid.uuid4())
+        trace_id = trace_id or time.time_ns()
+        if self.print_trace_id:
+            print(f"https://app.datadoghq.eu/logs?query=%40correlation_id%3A{trace_id}")
         search_id = self.search_id or (self._search_task.search_task_id if self._search_task is not None else None)
-        with MDC(trace_id=trace_id, search_id=search_id):
+        with MDC(correlation_id=trace_id, search_id=search_id):
             self.dump_input(trace_id, X)
             if len(args) > 0:
                 msg = f"WARNING: Unsupported positional arguments for transform: {args}"
@@ -904,10 +906,10 @@ class FeaturesEnricher(TransformerMixin):
             Dataframe with metrics calculated on train and validation datasets.
         """
 
-        trace_id = trace_id or str(uuid.uuid4())
+        trace_id = trace_id or time.time_ns()
         start_time = time.time()
         search_id = self.search_id or (self._search_task.search_task_id if self._search_task is not None else None)
-        with MDC(trace_id=trace_id, search_id=search_id):
+        with MDC(correlation_id=trace_id, search_id=search_id):
             self.logger.info("Start calculate metrics")
             if len(args) > 0:
                 msg = f"WARNING: Unsupported positional arguments for calculate_metrics: {args}"
@@ -1415,13 +1417,11 @@ class FeaturesEnricher(TransformerMixin):
         # Find latest eval set or earliest if all eval sets are before train set
         date_column = self._get_date_column(search_keys)
 
-        date_converter = DateTimeSearchKeyConverter(
+        date_converter = DateTimeConverter(
             date_column, self.date_format, self.logger, self.bundle, generate_cyclical_features=False
         )
 
-        X = date_converter.convert(X)
-
-        x_date = X[date_column].dropna()
+        x_date = date_converter.to_date_ms(X).dropna()
         if len(x_date) == 0:
             self.logger.warning("Empty date column in X")
             return []
@@ -1434,8 +1434,7 @@ class FeaturesEnricher(TransformerMixin):
             if date_column not in eval_x.columns:
                 self.logger.warning(f"Date column not found in eval_set {i + 1}")
                 continue
-            eval_x = date_converter.convert(eval_x)
-            eval_x_date = eval_x[date_column].dropna()
+            eval_x_date = date_converter.to_date_ms(eval_x).dropna()
             if len(eval_x_date) < 1000:
                 self.logger.warning(f"Eval_set {i} has less than 1000 rows. It will be ignored for stability check")
                 continue
@@ -1472,8 +1471,7 @@ class FeaturesEnricher(TransformerMixin):
         )
         checking_eval_set_df = checking_eval_set_df.copy()
 
-        checking_eval_set_df[date_column] = eval_set_dates[selected_eval_set_idx]
-        checking_eval_set_df = date_converter.convert(checking_eval_set_df)
+        checking_eval_set_df[date_column] = date_converter.to_date_ms(eval_set_dates[selected_eval_set_idx].to_frame())
 
         psi_values_sparse = calculate_sparsity_psi(
             checking_eval_set_df, cat_features, date_column, self.logger, model_task_type
@@ -1481,7 +1479,11 @@ class FeaturesEnricher(TransformerMixin):
 
         self.logger.info(f"PSI values by sparsity: {psi_values_sparse}")
 
-        unstable_by_sparsity = [feature for feature, psi in psi_values_sparse.items() if psi > stability_threshold]
+        unstable_by_sparsity = [
+            feature
+            for feature, psi in psi_values_sparse.items()
+            if psi > stability_threshold
+        ]
         if unstable_by_sparsity:
             self.logger.info(f"Unstable by sparsity features ({stability_threshold}): {sorted(unstable_by_sparsity)}")
 
@@ -1491,7 +1493,11 @@ class FeaturesEnricher(TransformerMixin):
 
         self.logger.info(f"PSI values by value: {psi_values}")
 
-        unstable_by_value = [feature for feature, psi in psi_values.items() if psi > stability_threshold]
+        unstable_by_value = [
+            feature
+            for feature, psi in psi_values.items()
+            if psi > stability_threshold
+        ]
         if unstable_by_value:
             self.logger.info(f"Unstable by value features ({stability_threshold}): {sorted(unstable_by_value)}")
 
@@ -1745,9 +1751,11 @@ class FeaturesEnricher(TransformerMixin):
             not in (
                 excluding_search_keys
                 + list(self.fit_dropped_features)
-                + [DateTimeSearchKeyConverter.DATETIME_COL, SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID]
+                + [DateTimeConverter.DATETIME_COL, SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID]
             )
         ]
+        if self.baseline_score_column is not None and self.baseline_score_column not in client_features:
+            client_features.append(self.baseline_score_column)
         self.logger.info(f"Client features column on prepare data for metrics: {client_features}")
 
         selected_enriched_features = [c for c in self.feature_names_ if c not in client_features]
@@ -1995,7 +2003,7 @@ class FeaturesEnricher(TransformerMixin):
         date_column = self._get_date_column(search_keys)
         generated_features = []
         if date_column is not None:
-            converter = DateTimeSearchKeyConverter(
+            converter = DateTimeConverter(
                 date_column,
                 self.date_format,
                 self.logger,
@@ -2004,6 +2012,7 @@ class FeaturesEnricher(TransformerMixin):
             )
             # Leave original date column values
             df_with_date_features = converter.convert(df, keep_time=True)
+            # TODO check if this is correct
             df_with_date_features[date_column] = df[date_column]
             df = df_with_date_features
             generated_features = converter.generated_features
@@ -2035,8 +2044,8 @@ class FeaturesEnricher(TransformerMixin):
         # Sample after sorting by system_record_id for idempotency
         df.sort_values(by=SYSTEM_RECORD_ID, inplace=True)
 
-        if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
-            df = df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL)
+        if DateTimeConverter.DATETIME_COL in df.columns:
+            df = df.drop(columns=DateTimeConverter.DATETIME_COL)
 
         df = df.rename(columns=columns_renaming)
         generated_features = [columns_renaming.get(c, c) for c in generated_features]
@@ -2388,7 +2397,7 @@ class FeaturesEnricher(TransformerMixin):
     def get_progress(self, trace_id: str | None = None, search_task: SearchTask | None = None) -> SearchProgress:
         search_task = search_task or self._search_task
         if search_task is not None:
-            trace_id = trace_id or uuid.uuid4()
+            trace_id = trace_id or time.time_ns()
             return search_task.get_progress(trace_id)
 
     def display_transactional_transform_api(self, only_online_sources=False):
@@ -2416,7 +2425,7 @@ class FeaturesEnricher(TransformerMixin):
                 return "12345678"
             return "test_value"
 
-        file_metadata = self._search_task.get_file_metadata(str(uuid.uuid4()))
+        file_metadata = self._search_task.get_file_metadata(time.time_ns())
 
         def get_column_meta(column_name: str) -> FileColumnMetadata:
             for c in file_metadata.columns:
@@ -2510,7 +2519,7 @@ if response.status_code == 200:
 
         start_time = time.time()
         search_id = self.search_id or (self._search_task.search_task_id if self._search_task is not None else None)
-        with MDC(trace_id=trace_id, search_id=search_id):
+        with MDC(correlation_id=trace_id, search_id=search_id):
             self.logger.info("Start transform")
 
             validated_X, validated_y, validated_eval_set = self._validate_train_eval(
@@ -2599,7 +2608,7 @@ if response.status_code == 200:
             generated_features = []
             date_column = self._get_date_column(search_keys)
             if date_column is not None:
-                converter = DateTimeSearchKeyConverter(
+                converter = DateTimeConverter(
                     date_column,
                     self.date_format,
                     self.logger,
@@ -2656,8 +2665,8 @@ if response.status_code == 200:
 
             # 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:
+            features_for_transform = self._search_task.get_features_for_transform()
+            if features_for_transform:
                 missing_features_for_transform = [
                     columns_renaming.get(f) or f for f in features_for_transform if f not in df.columns
                 ]
@@ -2668,7 +2677,10 @@ if response.status_code == 200:
                     raise ValidationError(
                         self.bundle.get("missing_features_for_transform").format(missing_features_for_transform)
                     )
-                runtime_parameters.properties["features_for_embeddings"] = ",".join(features_for_transform)
+            features_for_embeddings = self._search_task.get_features_for_embeddings()
+            if features_for_embeddings:
+                runtime_parameters.properties["features_for_embeddings"] = ",".join(features_for_embeddings)
+            features_for_transform = [f for f in features_for_transform if f not in search_keys.keys()]
 
             columns_for_system_record_id = sorted(list(search_keys.keys()) + features_for_transform)
 
@@ -2729,8 +2741,17 @@ if response.status_code == 200:
                 )
                 df = converter.convert(df)
 
+            date_features = []
+            for col in features_for_transform:
+                if DateTimeConverter(col).is_datetime(df):
+                    df[col] = DateTimeConverter(col).to_date_string(df)
+                    date_features.append(col)
+
             meaning_types = {}
-            meaning_types.update({col: FileColumnMeaningType.FEATURE for col in features_for_transform})
+            meaning_types.update(
+                {col: FileColumnMeaningType.FEATURE for col in features_for_transform if col not in date_features}
+            )
+            meaning_types.update({col: FileColumnMeaningType.DATE_FEATURE for col in date_features})
             meaning_types.update({col: key.value for col, key in search_keys.items()})
 
             features_not_to_pass.extend(
@@ -2743,8 +2764,8 @@ if response.status_code == 200:
                 ]
             )
 
-            if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
-                df = df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL)
+            if DateTimeConverter.DATETIME_COL in df.columns:
+                df = df.drop(columns=DateTimeConverter.DATETIME_COL)
 
             # search keys might be changed after explode
             columns_for_system_record_id = sorted(list(search_keys.keys()) + features_for_transform)
@@ -2926,6 +2947,7 @@ if response.status_code == 200:
                 or c in self.search_keys
                 or c in (self.id_columns or [])
                 or c in [EVAL_SET_INDEX, TARGET]  # transform for metrics calculation
+                or c == self.baseline_score_column
             ]
         else:
             selected_input_columns = []
@@ -3124,7 +3146,7 @@ if response.status_code == 200:
         self.fit_generated_features = []
 
         if has_date:
-            converter = DateTimeSearchKeyConverter(
+            converter = DateTimeConverter(
                 maybe_date_column,
                 self.date_format,
                 self.logger,
@@ -3177,8 +3199,8 @@ if response.status_code == 200:
             self.TARGET_NAME,
             EVAL_SET_INDEX,
         ] + list(self.fit_search_keys.keys())
-        if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
-            non_feature_columns.append(DateTimeSearchKeyConverter.DATETIME_COL)
+        if DateTimeConverter.DATETIME_COL in df.columns:
+            non_feature_columns.append(DateTimeConverter.DATETIME_COL)
 
         features_columns = [c for c in df.columns if c not in non_feature_columns]
 
@@ -3265,15 +3287,27 @@ if response.status_code == 200:
             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)
+        if DateTimeConverter.DATETIME_COL in df.columns:
+            non_feature_columns.append(DateTimeConverter.DATETIME_COL)
 
         features_columns = [c for c in df.columns if c not in non_feature_columns]
 
+        # find date features
+        date_features = []
+        for col in features_columns:
+            if DateTimeConverter(col).is_datetime(df):
+                df[col] = DateTimeConverter(col).to_date_string(df)
+                date_features.append(col)
+
         meaning_types = {
             **{col: key.value for col, key in self.fit_search_keys.items()},
-            **{str(c): FileColumnMeaningType.FEATURE for c in df.columns if c not in non_feature_columns},
+            **{
+                str(c): FileColumnMeaningType.FEATURE
+                for c in df.columns
+                if c not in non_feature_columns and c not in date_features
+            },
         }
+        meaning_types.update({col: FileColumnMeaningType.DATE_FEATURE for col in date_features})
         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:
@@ -3294,8 +3328,8 @@ if response.status_code == 200:
             self.bundle,
         )
 
-        if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
-            df = df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL)
+        if DateTimeConverter.DATETIME_COL in df.columns:
+            df = df.drop(columns=DateTimeConverter.DATETIME_COL)
 
         meaning_types[SYSTEM_RECORD_ID] = FileColumnMeaningType.SYSTEM_RECORD_ID
 
@@ -3332,7 +3366,9 @@ if response.status_code == 200:
         dataset.columns_renaming = self.fit_columns_renaming
 
         self.passed_features = [
-            column for column, meaning_type in meaning_types.items() if meaning_type == FileColumnMeaningType.FEATURE
+            column
+            for column, meaning_type in meaning_types.items()
+            if meaning_type in [FileColumnMeaningType.FEATURE, FileColumnMeaningType.DATE_FEATURE]
         ]
 
         self._search_task = dataset.search(
@@ -3860,8 +3896,8 @@ if response.status_code == 200:
             X = Xy.drop(columns=TARGET)
             y = Xy[TARGET].copy()
 
-        if DateTimeSearchKeyConverter.DATETIME_COL in X.columns:
-            X.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL, inplace=True)
+        if DateTimeConverter.DATETIME_COL in X.columns:
+            X.drop(columns=DateTimeConverter.DATETIME_COL, inplace=True)
 
         return X, y
 
@@ -3871,8 +3907,8 @@ if response.status_code == 200:
         X: pd.DataFrame, y: pd.Series, search_keys: dict[str, SearchKey], cv: CVType | None
     ) -> tuple[pd.DataFrame, pd.Series]:
         if cv not in [CVType.time_series, CVType.blocked_time_series]:
-            if DateTimeSearchKeyConverter.DATETIME_COL in X.columns:
-                date_column = DateTimeSearchKeyConverter.DATETIME_COL
+            if DateTimeConverter.DATETIME_COL in X.columns:
+                date_column = DateTimeConverter.DATETIME_COL
             else:
                 date_column = FeaturesEnricher._get_date_column(search_keys)
             sort_columns = [date_column] if date_column is not None else []
@@ -3900,8 +3936,8 @@ if response.status_code == 200:
 
             y = Xy[TARGET].copy()
 
-        if DateTimeSearchKeyConverter.DATETIME_COL in X.columns:
-            X.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL, inplace=True)
+        if DateTimeConverter.DATETIME_COL in X.columns:
+            X.drop(columns=DateTimeConverter.DATETIME_COL, inplace=True)
 
         return X, y
 
@@ -3980,12 +4016,10 @@ if response.status_code == 200:
             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, generate_cyclical_features=False
-                )
-                converted_X = date_converter.convert(X)
-                min_date = converted_X[maybe_date_col].min()
-                max_date = converted_X[maybe_date_col].max()
+                date_converter = DateTimeConverter(maybe_date_col, self.date_format, generate_cyclical_features=False)
+                date_col_values = date_converter.to_date_ms(X)
+                min_date = date_col_values.min()
+                max_date = date_col_values.max()
                 self.logger.info(f"Dates interval is ({min_date}, {max_date})")
 
         except Exception:
@@ -4022,7 +4056,7 @@ if response.status_code == 200:
                 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, generate_cyclical_features=False)
+            converter = DateTimeConverter(FeaturesEnricher.CURRENT_DATE, generate_cyclical_features=False)
             df = converter.convert(df)
         return df
 
@@ -4153,8 +4187,8 @@ if response.status_code == 200:
             "__target",
             ENTITY_SYSTEM_RECORD_ID,
         ]
-        if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
-            date_column = DateTimeSearchKeyConverter.DATETIME_COL
+        if DateTimeConverter.DATETIME_COL in df.columns:
+            date_column = DateTimeConverter.DATETIME_COL
             sort_exclude_columns.append(FeaturesEnricher._get_date_column(search_keys))
         else:
             date_column = FeaturesEnricher._get_date_column(search_keys)
@@ -4954,7 +4988,7 @@ if response.status_code == 200:
         eval_set: tuple | None = None,
     ):
         def dump_task(X_, y_, eval_set_):
-            with MDC(trace_id=trace_id):
+            with MDC(correlation_id=trace_id):
                 try:
                     if isinstance(X_, pd.Series):
                         X_ = X_.to_frame()

upgini/metadata.py

@@ -36,6 +36,7 @@ class FileColumnMeaningType(Enum):
     SCORE = "SCORE"
     TARGET = "TARGET"
     FEATURE = "FEATURE"
+    DATE_FEATURE = "DATE_FEATURE"
     CUSTOM_KEY = "CUSTOM_KEY"
     COUNTRY = "COUNTRY"
     POSTAL_CODE = "POSTAL_CODE"

upgini/normalizer/normalize_utils.py

@@ -25,7 +25,7 @@ from upgini.metadata import (
 from upgini.resource_bundle import ResourceBundle, get_custom_bundle
 from upgini.utils import find_numbers_with_decimal_comma
 from upgini.utils.country_utils import CountrySearchKeyConverter
-from upgini.utils.datetime_utils import DateTimeSearchKeyConverter
+from upgini.utils.datetime_utils import DateTimeConverter
 from upgini.utils.ip_utils import IpSearchKeyConverter
 from upgini.utils.phone_utils import PhoneSearchKeyConverter
 from upgini.utils.postal_code_utils import PostalCodeSearchKeyConverter
@@ -89,7 +89,7 @@ class Normalizer:
                 SYSTEM_RECORD_ID,
                 ENTITY_SYSTEM_RECORD_ID,
                 SEARCH_KEY_UNNEST,
-                DateTimeSearchKeyConverter.DATETIME_COL,
+                DateTimeConverter.DATETIME_COL,
             ]:
                 self.columns_renaming[column] = column
                 new_columns.append(column)

upgini/search_task.py

@@ -165,10 +165,21 @@ class SearchTask:
 
         return list(zero_hit_search_keys)
 
-    def get_features_for_transform(self) -> Optional[List[str]]:
+    def get_features_for_embeddings(self) -> Optional[List[str]]:
         if self.provider_metadata_v2 is None:
             return None
 
+        features_for_transform = set()
+        for meta in self.provider_metadata_v2:
+            if meta.features_used_for_embeddings is not None:
+                features_for_transform.update(meta.features_used_for_embeddings)
+
+        return list(features_for_transform)
+
+    def get_features_for_transform(self) -> List[str]:
+        if self.provider_metadata_v2 is None:
+            return []
+
         features_for_transform = set()
         for meta in self.provider_metadata_v2:
             if meta.features_used_for_embeddings is not None:

upgini/utils/datetime_utils.py

@@ -30,7 +30,7 @@ DATE_FORMATS = [
 DATETIME_PATTERN = r"^[\d\s\.\-:T/+]+$"
 
 
-class DateTimeSearchKeyConverter:
+class DateTimeConverter:
     DATETIME_COL = "_date_time"
     # MIN_SUPPORTED_DATE_TS = datetime.datetime(1999, 12, 31)  # 946684800000  # 2000-01-01
     MIN_SUPPORTED_DATE_TS = pd.to_datetime(datetime.datetime(1999, 12, 31)).tz_localize(None)
@@ -73,41 +73,99 @@ class DateTimeSearchKeyConverter:
         except Exception:
             return None
 
-    def convert(self, df: pd.DataFrame, keep_time=False) -> pd.DataFrame:
-        if len(df) == 0:
-            return df
+    def is_datetime(self, df: pd.DataFrame) -> bool:
+        if len(df) == 0 or df[self.date_column].isna().all():
+            return False
+
+        if pd.api.types.is_datetime64_any_dtype(df[self.date_column]):
+            return True
+
+        parsed = self.parse_datetime(df, raise_errors=False)
+        return parsed is not None and not parsed.isna().all()
 
+    def parse_datetime(self, df: pd.DataFrame, raise_errors=True) -> pd.Series | None:
         df = df.copy()
-        if df[self.date_column].apply(lambda x: isinstance(x, datetime.datetime)).all():
-            df[self.date_column] = df[self.date_column].apply(lambda x: x.replace(tzinfo=None))
-        elif isinstance(df[self.date_column].values[0], datetime.date):
-            df[self.date_column] = pd.to_datetime(df[self.date_column], errors="coerce")
-        elif isinstance(df[self.date_column].dtype, pd.PeriodDtype):
-            df[self.date_column] = df[self.date_column].dt.to_timestamp()
-        elif is_numeric_dtype(df[self.date_column]):
-            # 315532801 - 2524608001    - seconds
-            # 315532801000 - 2524608001000 - milliseconds
-            # 315532801000000 - 2524608001000000 - microseconds
-            # 315532801000000000 - 2524608001000000000 - nanoseconds
-            if df[self.date_column].apply(lambda x: 10**16 < x).all():
-                df[self.date_column] = pd.to_datetime(df[self.date_column], unit="ns")
-            elif df[self.date_column].apply(lambda x: 10**14 < x < 10**16).all():
-                df[self.date_column] = pd.to_datetime(df[self.date_column], unit="us")
-            elif df[self.date_column].apply(lambda x: 10**11 < x < 10**14).all():
-                df[self.date_column] = pd.to_datetime(df[self.date_column], unit="ms")
-            elif df[self.date_column].apply(lambda x: 0 < x < 10**11).all():
-                df[self.date_column] = pd.to_datetime(df[self.date_column], unit="s")
+        if len(df) == 0 or df[self.date_column].isna().all():
+            return None
+
+        try:
+            if df[self.date_column].apply(lambda x: isinstance(x, datetime.datetime)).all():
+                parsed_datetime = df[self.date_column].apply(lambda x: x.replace(tzinfo=None))
+            elif isinstance(df[self.date_column].dropna().values[0], datetime.date):
+                parsed_datetime = pd.to_datetime(df[self.date_column], errors="coerce")
+            elif isinstance(df[self.date_column].dtype, pd.PeriodDtype):
+                parsed_datetime = df[self.date_column].dt.to_timestamp()
+            elif is_numeric_dtype(df[self.date_column]):
+                # 315532801 - 2524608001    - seconds
+                # 315532801000 - 2524608001000 - milliseconds
+                # 315532801000000 - 2524608001000000 - microseconds
+                # 315532801000000000 - 2524608001000000000 - nanoseconds
+                if df[self.date_column].apply(lambda x: 10**16 < x).all():
+                    parsed_datetime = pd.to_datetime(df[self.date_column], unit="ns")
+                elif df[self.date_column].apply(lambda x: 10**14 < x < 10**16).all():
+                    parsed_datetime = pd.to_datetime(df[self.date_column], unit="us")
+                elif df[self.date_column].apply(lambda x: 10**11 < x < 10**14).all():
+                    parsed_datetime = pd.to_datetime(df[self.date_column], unit="ms")
+                elif df[self.date_column].apply(lambda x: 10**8 < x < 10**11).all():
+                    parsed_datetime = pd.to_datetime(df[self.date_column], unit="s")
+                else:
+                    msg = self.bundle.get("unsupported_date_type").format(self.date_column)
+                    if raise_errors:
+                        raise ValidationError(msg)
+                    else:
+                        return None
+            else:
+                df[self.date_column] = df[self.date_column].astype("string").apply(self.clean_date)
+                parsed_datetime = self.parse_string_date(df, raise_errors)
+            parsed_datetime = parsed_datetime.dt.tz_localize(None)
+            return parsed_datetime
+        except Exception as e:
+            if raise_errors:
+                raise ValidationError(e)
             else:
-                msg = self.bundle.get("unsupported_date_type").format(self.date_column)
-                raise ValidationError(msg)
+                return None
+
+    def to_date_string(self, df: pd.DataFrame) -> pd.Series:
+        parsed_datetime = self.parse_datetime(df)
+        if parsed_datetime is None:
+            return df[self.date_column]
+        return parsed_datetime.dt.strftime("%Y-%m-%d")
+
+    def to_date_ms(self, df: pd.DataFrame) -> pd.Series:
+        parsed_datetime = self.parse_datetime(df)
+        if parsed_datetime is None:
+            return df[self.date_column]
+        return self.convert_datetime_to_date_ms(parsed_datetime)
+
+    def convert_datetime_to_datetime_ms(self, date_col: pd.Series) -> pd.Series:
+        if date_col.dt.unit == "ns":
+            date_col = date_col.astype(np.int64) // 1_000_000
+        elif date_col.dt.unit == "us":
+            date_col = date_col.astype(np.int64) // 1_000
+        elif date_col.dt.unit == "ms":
+            date_col = date_col.astype(np.int64)
+        elif date_col.dt.unit == "s":
+            date_col = date_col.astype(np.int64) * 1_000
         else:
-            df[self.date_column] = df[self.date_column].astype("string").apply(self.clean_date)
-            df[self.date_column] = self.parse_date(df)
+            raise ValueError(f"Unsupported date unit: {date_col.dt.unit}")
+
+        return date_col.apply(self._int_to_opt).astype("Int64")
+
+    def convert_datetime_to_date_ms(self, date_col: pd.Series) -> pd.Series:
+        date_col = date_col.dt.floor("D")
+        return self.convert_datetime_to_datetime_ms(date_col)
+
+    def convert(self, df: pd.DataFrame, keep_time=False) -> pd.DataFrame:
+        df = df.copy()
+        parsed_datetime = self.parse_datetime(df)
+        if parsed_datetime is None:
+            return df
+
+        df[self.date_column] = parsed_datetime
 
         # If column with date is datetime then extract seconds of the day and minute of the hour
         # as additional features
         seconds = "datetime_seconds"
-        df[self.date_column] = df[self.date_column].dt.tz_localize(None)
 
         df = self.clean_old_dates(df)
 
@@ -182,21 +240,22 @@ class DateTimeSearchKeyConverter:
             df.drop(columns=seconds, inplace=True)
 
         if keep_time:
-            df[self.DATETIME_COL] = df[self.date_column].astype(np.int64) // 1_000_000
-            df[self.DATETIME_COL] = df[self.DATETIME_COL].apply(self._int_to_opt).astype("Int64")
-        df[self.date_column] = df[self.date_column].dt.floor("D").astype(np.int64) // 1_000_000
-        df[self.date_column] = df[self.date_column].apply(self._int_to_opt).astype("Int64")
+            df[self.DATETIME_COL] = self.convert_datetime_to_datetime_ms(df[self.date_column])
+        df[self.date_column] = self.convert_datetime_to_date_ms(df[self.date_column])
 
         self.logger.info(f"Date after convertion to timestamp: {df[self.date_column]}")
 
         return df
 
-    def parse_date(self, df: pd.DataFrame):
+    def parse_string_date(self, df: pd.DataFrame, raise_errors=True) -> pd.Series | None:
         if self.date_format is not None:
             try:
                 return pd.to_datetime(df[self.date_column], format=self.date_format)
             except ValueError as e:
-                raise ValidationError(e)
+                if raise_errors:
+                    raise ValidationError(e)
+                else:
+                    return None
         else:
             for date_format in DATE_FORMATS:
                 try:
@@ -204,9 +263,17 @@ class DateTimeSearchKeyConverter:
                 except ValueError:
                     pass
             try:
-                return pd.to_datetime(df[self.date_column])
+                # Suppress warning for intentional fallback to dateutil parsing
+                import warnings
+
+                with warnings.catch_warnings():
+                    warnings.filterwarnings("ignore", message="Could not infer format")
+                    return pd.to_datetime(df[self.date_column])
             except ValueError:
-                raise ValidationError(self.bundle.get("invalid_date_format").format(self.date_column))
+                if raise_errors:
+                    raise ValidationError(self.bundle.get("invalid_date_format").format(self.date_column))
+                else:
+                    return None
 
     def clean_old_dates(self, df: pd.DataFrame) -> pd.DataFrame:
         condition = df[self.date_column] <= self.MIN_SUPPORTED_DATE_TS

upgini/utils/deduplicate_utils.py

@@ -14,7 +14,7 @@ from upgini.metadata import (
     SearchKey,
 )
 from upgini.resource_bundle import ResourceBundle, get_custom_bundle
-from upgini.utils.datetime_utils import DateTimeSearchKeyConverter
+from upgini.utils.datetime_utils import DateTimeConverter
 from upgini.utils.target_utils import define_task
 
 
@@ -104,7 +104,7 @@ def remove_fintech_duplicates(
         sub_df = pd.merge(sub_df, nonunique_target_rows, on=personal_cols)
 
         # Convert date columns for further checks
-        sub_df = DateTimeSearchKeyConverter(
+        sub_df = DateTimeConverter(
             date_col, date_format=date_format, logger=logger, bundle=bundle, generate_cyclical_features=False
         ).convert(sub_df)
         grouped_by_personal_cols = sub_df.groupby(personal_cols, group_keys=False)

{upgini-1.2.124.dist-info → upgini-1.2.125.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.124
+Version: 1.2.125
 Summary: Intelligent data search & enrichment for Machine Learning
 Project-URL: Bug Reports, https://github.com/upgini/upgini/issues
 Project-URL: Homepage, https://upgini.com/
@@ -30,6 +30,7 @@ Requires-Dist: ipywidgets>=8.1.0
 Requires-Dist: jarowinkler>=2.0.0
 Requires-Dist: levenshtein>=0.25.1
 Requires-Dist: lightgbm>=4.6.0
+Requires-Dist: more-itertools==10.7.0
 Requires-Dist: numpy<3.0.0,>=1.19.0
 Requires-Dist: pandas<3.0.0,>=1.1.0
 Requires-Dist: psutil>=5.9.0

{upgini-1.2.124.dist-info → upgini-1.2.125.dist-info}/RECORD

@@ -1,20 +1,20 @@
-upgini/__about__.py,sha256=BVLX26XWHqzeaXiLce8Wr2FrayZ3daxZy_Sdqvu_jrI,24
+upgini/__about__.py,sha256=khvL6Ma3KHnaaXtUCPR9kKBJFG5qg7emKoKVlrbEt0k,24
 upgini/__init__.py,sha256=LXSfTNU0HnlOkE69VCxkgIKDhWP-JFo_eBQ71OxTr5Y,261
 upgini/ads.py,sha256=nvuRxRx5MHDMgPr9SiU-fsqRdFaBv8p4_v1oqiysKpc,2714
-upgini/dataset.py,sha256=pQ8JQe0cdygD-W9GefJmfE6bnj4EYzXsjlgWdIS9nS8,31578
+upgini/dataset.py,sha256=Nm2ZmwyQqvTnymYpGUwyJWy7y2ebXlHMyYmGeGcyA_s,31652
 upgini/errors.py,sha256=2b_Wbo0OYhLUbrZqdLIx5jBnAsiD1Mcenh-VjR4HCTw,950
-upgini/features_enricher.py,sha256=unkzwcMx2kcCqSYzjOP3JzcCsIYMZs30nc0mviOLw5Y,231791
+upgini/features_enricher.py,sha256=tmKeERG2b0YfJ47g-UXQQ3S-9tyagwUOhI4oqN3kG2w,233058
 upgini/http.py,sha256=-J_wOpnwVnT0ebPC6sOs6fN3AWtCD0LJLu6nlYmxaqk,44348
-upgini/metadata.py,sha256=1j2LkkPmFAVaVmT__FT3nYNecf9NMpwY3YE_HN_YcrY,12541
+upgini/metadata.py,sha256=CL9bFytdUZlbQYtTgNgAkt_sxO9klARQtULDBgb2Hlg,12575
 upgini/metrics.py,sha256=KCPE_apPN-9BIdv6GqASbJVaB_gBcy8wzNApAcyaGo4,46020
-upgini/search_task.py,sha256=SAiUd1AytbA2Q6PSnnztr7oTRKpud1wQZ5YtKjsmQHU,18256
+upgini/search_task.py,sha256=5mL_qV5mVtDkIumM9xCOgfa9Lc2B8mxJ1qI21iaScnQ,18656
 upgini/spinner.py,sha256=4iMd-eIe_BnkqFEMIliULTbj6rNI2HkN_VJ4qYe0cUc,1118
 upgini/version_validator.py,sha256=DvbaAvuYFoJqYt0fitpsk6Xcv-H1BYDJYHUMxaKSH_Y,1509
 upgini/ads_management/__init__.py,sha256=qzyisOToVRP-tquAJD1PblZhNtMrOB8FiyF9JvfkvgE,50
 upgini/ads_management/ads_manager.py,sha256=igVbN2jz80Umb2BUJixmJVj-zx8unoKpecVo-R-nGdw,2648
 upgini/autofe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 upgini/autofe/all_operators.py,sha256=rdjF5eaE4bC6Q4eu_el5Z7ekYt8DjOFermz2bePPbUc,333
-upgini/autofe/binary.py,sha256=oOEECc4nRzZN2tYaiqx8F2XHnfWpk1bVvb7ZkZJ0lO8,7709
+upgini/autofe/binary.py,sha256=o3TQuP3EnECAVIeToGczu4yJ4vX7BJ2iSCN9Ra1SZJI,7829
 upgini/autofe/date.py,sha256=RvexgrL1_6ISYPVrl9HUQmPgpVSGQsTNv8YhNQWs-5M,11329
 upgini/autofe/feature.py,sha256=W9sZHdz5Vi0H_oPyY5saZAPjyd5wunpULnCqrGLpQc4,16879
 upgini/autofe/groupby.py,sha256=IYmQV9uoCdRcpkeWZj_kI3ObzoNCNx3ff3h8sTL01tk,3603
@@ -35,7 +35,7 @@ upgini/data_source/data_source_publisher.py,sha256=qXQUYErhCmkWHm2FWgTL0FYZ2aJbx
 upgini/mdc/__init__.py,sha256=iHJlXQg6xRM1-ZOUtaPSJqw5SpQDszvxp4LyqviNLIQ,1027
 upgini/mdc/context.py,sha256=3u1B-jXt7tXEvNcV3qmR9SDCseudnY7KYsLclBdwVLk,1405
 upgini/normalizer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-upgini/normalizer/normalize_utils.py,sha256=mDh2mBW3aQMB4EFP2aHbf2dGMVkOcWnp4sKKvKDBh8w,8511
+upgini/normalizer/normalize_utils.py,sha256=w9f_9udrwqbhXgFMTs2keuce-6X_j6h3D7EdNo_2X7g,8493
 upgini/resource_bundle/__init__.py,sha256=S5F2G47pnJd2LDpmFsjDqEwiKkP8Hm-hcseDbMka6Ko,8345
 upgini/resource_bundle/exceptions.py,sha256=5fRvx0_vWdE1-7HcSgF0tckB4A9AKyf5RiinZkInTsI,621
 upgini/resource_bundle/strings.properties,sha256=KcXm1Nl6c3zswL91tIbG0DjuuNpzxUdCg1cY9f2-9cg,29283
@@ -52,8 +52,8 @@ upgini/utils/config.py,sha256=zFdnjchykfp_1Tm3Qep7phLzXBpXIOzr2tIuXchRBLw,1754
 upgini/utils/country_utils.py,sha256=lY-eXWwFVegdVENFttbvLcgGDjFO17Sex8hd2PyJaRk,6937
 upgini/utils/custom_loss_utils.py,sha256=kieNZYBYZm5ZGBltF1F_jOSF4ea6C29rYuCyiDcqVNY,3857
 upgini/utils/cv_utils.py,sha256=w6FQb9nO8BWDx88EF83NpjPLarK4eR4ia0Wg0kLBJC4,3525
-upgini/utils/datetime_utils.py,sha256=UL1ernnawW0LV9mPDpCIc6sFy0HUhFscWVNwfH4V7rI,14366
-upgini/utils/deduplicate_utils.py,sha256=oZEiZeN-A92zwAPysV4OP9hO-niC2RLt-Dhc_hynBTU,11273
+upgini/utils/datetime_utils.py,sha256=l85UzSQLhtMeI2G6m-m8y8bCColCLSXNHb2-G6fKpLM,16988
+upgini/utils/deduplicate_utils.py,sha256=6czbn1q0p-lOmrNvbAzueBpDHmfIP4TfV4poWqbjX5w,11255
 upgini/utils/display_utils.py,sha256=uSG3JwpwCIgRJXsp-8ktuJ0Dh-WFti7IrRLMUfHfoDc,11973
 upgini/utils/email_utils.py,sha256=pZ2vCfNxLIPUhxr0-OlABNXm12jjU44isBk8kGmqQzA,5277
 upgini/utils/fallback_progress_bar.py,sha256=PDaKb8dYpVZaWMroNcOHsTc3pSjgi9mOm0--cOFTwJ0,1074
@@ -74,7 +74,7 @@ upgini/utils/target_utils.py,sha256=GCPn4QeJ83JJ_vyBJ3IhY5fyIRkLC9q9BE59S2FRO1I,
 upgini/utils/track_info.py,sha256=G5Lu1xxakg2_TQjKZk4b5SvrHsATTXNVV3NbvWtT8k8,5663
 upgini/utils/ts_utils.py,sha256=26vhC0pN7vLXK6R09EEkMK3Lwb9IVPH7LRdqFIQ3kPs,1383
 upgini/utils/warning_counter.py,sha256=-GRY8EUggEBKODPSuXAkHn9KnEQwAORC0mmz_tim-PM,254
-upgini-1.2.124.dist-info/METADATA,sha256=4oYOnQTr7xUDah86SaJL141wp3linADEiUVcY0X60Vk,50743
-upgini-1.2.124.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-upgini-1.2.124.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
-upgini-1.2.124.dist-info/RECORD,,
+upgini-1.2.125.dist-info/METADATA,sha256=CAoP8m15syLZEVmnYuUjUMI1Jo-XvMCGhz-CZnRYwy4,50781
+upgini-1.2.125.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+upgini-1.2.125.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
+upgini-1.2.125.dist-info/RECORD,,