upgini 1.2.36a1__tar.gz → 1.2.38__tar.gz

{upgini-1.2.36a1 → upgini-1.2.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.36a1
+Version: 1.2.38
 Summary: Intelligent data search & enrichment for Machine Learning
 Project-URL: Bug Reports, https://github.com/upgini/upgini/issues
 Project-URL: Homepage, https://upgini.com/

upgini-1.2.38/src/upgini/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.2.38"

{upgini-1.2.36a1 → upgini-1.2.38}/src/upgini/dataset.py

@@ -22,6 +22,7 @@ from upgini.metadata import (
     EVAL_SET_INDEX,
     SYSTEM_RECORD_ID,
     TARGET,
+    CVType,
     DataType,
     FeaturesFilter,
     FileColumnMeaningType,
@@ -32,11 +33,12 @@ from upgini.metadata import (
     NumericInterval,
     RuntimeParameters,
     SearchCustomization,
+    SearchKey,
 )
 from upgini.resource_bundle import ResourceBundle, get_custom_bundle
 from upgini.search_task import SearchTask
 from upgini.utils.email_utils import EmailSearchKeyConverter
-from upgini.utils.target_utils import balance_undersample, balance_undersample_forced
+from upgini.utils.target_utils import balance_undersample, balance_undersample_forced, balance_undersample_time_series
 
 try:
     from upgini.utils.progress_bar import CustomProgressBar as ProgressBar
@@ -74,6 +76,8 @@ class Dataset:  # (pd.DataFrame):
         search_keys: Optional[List[Tuple[str, ...]]] = None,
         unnest_search_keys: Optional[Dict[str, str]] = None,
         model_task_type: Optional[ModelTaskType] = None,
+        cv_type: Optional[CVType] = None,
+        id_columns: Optional[List[str]] = None,
         random_state: Optional[int] = None,
         rest_client: Optional[_RestClient] = None,
         logger: Optional[logging.Logger] = None,
@@ -104,6 +108,7 @@ class Dataset:  # (pd.DataFrame):
 
         self.dataset_name = dataset_name
         self.task_type = model_task_type
+        self.cv_type = cv_type
         self.description = description
         self.meaning_types = meaning_types
         self.search_keys = search_keys
@@ -116,6 +121,7 @@ class Dataset:  # (pd.DataFrame):
         self.random_state = random_state
         self.columns_renaming: Dict[str, str] = {}
         self.imbalanced: bool = False
+        self.id_columns = id_columns
         if logger is not None:
             self.logger = logger
         else:
@@ -225,6 +231,8 @@ class Dataset:  # (pd.DataFrame):
                 df=self.data,
                 target_column=target_column,
                 task_type=self.task_type,
+                cv_type=self.cv_type,
+                id_columns=self.id_columns,
                 random_state=self.random_state,
                 sample_size=self.FORCE_SAMPLE_SIZE,
                 logger=self.logger,
@@ -297,7 +305,21 @@ class Dataset:  # (pd.DataFrame):
                 f"Etalon has size {len(self.data)} more than threshold {sample_threshold} "
                 f"and will be downsampled to {sample_rows}"
             )
-            resampled_data = self.data.sample(n=sample_rows, random_state=self.random_state)
+            if self.cv_type is not None and self.cv_type.is_time_series():
+                resampled_data = balance_undersample_time_series(
+                    df=self.data,
+                    id_columns=self.id_columns,
+                    date_column=next(
+                        k
+                        for k, v in self.meaning_types.items()
+                        if v in [FileColumnMeaningType.DATE, FileColumnMeaningType.DATETIME]
+                    ),
+                    sample_size=sample_rows,
+                    random_state=self.random_state,
+                    logger=self.logger,
+                )
+            else:
+                resampled_data = self.data.sample(n=sample_rows, random_state=self.random_state)
             self.data = resampled_data
             self.logger.info(f"Shape after threshold resampling: {self.data.shape}")
 

{upgini-1.2.36a1 → upgini-1.2.38}/src/upgini/features_enricher.py

@@ -237,6 +237,7 @@ class FeaturesEnricher(TransformerMixin):
         add_date_if_missing: bool = True,
         select_features: bool = False,
         disable_force_downsampling: bool = False,
+        id_columns: Optional[List[str]] = None,
         **kwargs,
     ):
         self.bundle = get_custom_bundle(custom_bundle_config)
@@ -277,9 +278,12 @@ class FeaturesEnricher(TransformerMixin):
         )
 
         validate_version(self.logger, self.__log_warning)
+
         self.search_keys = search_keys or {}
+        self.id_columns = id_columns
         self.country_code = country_code
         self.__validate_search_keys(search_keys, search_id)
+
         self.model_task_type = model_task_type
         self.endpoint = endpoint
         self._search_task: Optional[SearchTask] = None
@@ -928,6 +932,8 @@ class FeaturesEnricher(TransformerMixin):
                 cat_features, search_keys_for_metrics = self._get_client_cat_features(
                     estimator, validated_X, self.search_keys
                 )
+                search_keys_for_metrics.extend([c for c in self.id_columns or [] if c not in search_keys_for_metrics])
+                self.logger.info(f"Search keys for metrics: {search_keys_for_metrics}")
 
                 prepared_data = self._prepare_data_for_metrics(
                     trace_id=trace_id,
@@ -983,7 +989,7 @@ class FeaturesEnricher(TransformerMixin):
                 with Spinner():
                     self._check_train_and_eval_target_distribution(y_sorted, fitting_eval_set_dict)
 
-                    has_date = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME]) is not None
+                    has_date = self._get_date_column(search_keys) is not None
                     model_task_type = self.model_task_type or define_task(y_sorted, has_date, self.logger, silent=True)
 
                     wrapper = EstimatorWrapper.create(
@@ -1185,7 +1191,7 @@ class FeaturesEnricher(TransformerMixin):
                     )
 
                     uplift_col = self.bundle.get("quality_metrics_uplift_header")
-                    date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+                    date_column = self._get_date_column(search_keys)
                     if (
                         uplift_col in metrics_df.columns
                         and (metrics_df[uplift_col] < 0).any()
@@ -1354,7 +1360,7 @@ class FeaturesEnricher(TransformerMixin):
         groups = None
 
         if not isinstance(_cv, BaseCrossValidator):
-            date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+            date_column = self._get_date_column(search_keys)
             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
@@ -1443,9 +1449,13 @@ class FeaturesEnricher(TransformerMixin):
 
         excluding_search_keys = list(search_keys.keys())
         if search_keys_for_metrics is not None and len(search_keys_for_metrics) > 0:
+            excluded = set()
             for sk in excluding_search_keys:
                 if columns_renaming.get(sk) in search_keys_for_metrics:
-                    excluding_search_keys.remove(sk)
+                    excluded.add(sk)
+            excluding_search_keys = [sk for sk in excluding_search_keys if sk not in excluded]
+
+        self.logger.info(f"Excluding search keys: {excluding_search_keys}")
 
         client_features = [
             c
@@ -1667,7 +1677,7 @@ class FeaturesEnricher(TransformerMixin):
         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])
+        date_column = self._get_date_column(search_keys)
         generated_features = []
         if date_column is not None:
             converter = DateTimeSearchKeyConverter(date_column, self.date_format, self.logger, self.bundle)
@@ -1741,7 +1751,7 @@ class FeaturesEnricher(TransformerMixin):
         search_keys = self.fit_search_keys
 
         rows_to_drop = None
-        has_date = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME]) is not None
+        has_date = self._get_date_column(search_keys) 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
         )
@@ -1853,7 +1863,10 @@ class FeaturesEnricher(TransformerMixin):
                 df = balance_undersample_forced(
                     df=df,
                     target_column=TARGET,
+                    id_columns=self.id_columns,
+                    date_column=self._get_date_column(self.search_keys),
                     task_type=self.model_task_type,
+                    cv_type=self.cv,
                     random_state=self.random_state,
                     sample_size=Dataset.FORCE_SAMPLE_SIZE,
                     logger=self.logger,
@@ -1995,7 +2008,7 @@ class FeaturesEnricher(TransformerMixin):
             trace_id = trace_id or uuid.uuid4()
             return search_task.get_progress(trace_id)
 
-    def get_transactional_transform_api(self):
+    def get_transactional_transform_api(self, only_online_sources=False):
         if self.api_key is None:
             raise ValidationError(self.bundle.get("transactional_transform_unregistered"))
         if self._search_task is None:
@@ -2053,7 +2066,7 @@ class FeaturesEnricher(TransformerMixin):
         api_example = f"""curl 'https://search.upgini.com/online/api/http_inference_trigger?search_id={search_id}' \\
     -H 'Authorization: {self.api_key}' \\
     -H 'Content-Type: application/json' \\
-    -d '{{"search_keys": {keys}{features_section}}}'"""
+    -d '{{"search_keys": {keys}{features_section}, "only_online_sources": {str(only_online_sources).lower()}}}'"""
         return api_example
 
     def _get_copy_of_runtime_parameters(self) -> RuntimeParameters:
@@ -2097,13 +2110,15 @@ class FeaturesEnricher(TransformerMixin):
                 return None, {c: c for c in X.columns}, []
 
             features_meta = self._search_task.get_all_features_metadata_v2()
-            online_api_features = [fm.name for fm in features_meta if fm.from_online_api]
+            online_api_features = [fm.name for fm in features_meta if fm.from_online_api and fm.shap_value > 0]
             if len(online_api_features) > 0:
                 self.logger.warning(
                     f"There are important features for transform, that generated by online API: {online_api_features}"
                 )
-                # TODO
-                raise Exception("There are features selected that are paid. Contact support (sales@upgini.com)")
+                msg = self.bundle.get("online_api_features_transform").format(online_api_features)
+                self.logger.warning(msg)
+                print(msg)
+                print(self.get_transactional_transform_api(only_online_sources=True))
 
             if not metrics_calculation:
                 transform_usage = self.rest_client.get_current_transform_usage(trace_id)
@@ -2135,6 +2150,9 @@ class FeaturesEnricher(TransformerMixin):
                 validated_X = validated_X.drop(columns=columns_to_drop)
 
             search_keys = self.search_keys.copy()
+            if self.id_columns is not None and self.cv is not None and self.cv.is_time_series():
+                self.search_keys.update({col: SearchKey.CUSTOM_KEY for col in self.id_columns})
+                
             search_keys = self.__prepare_search_keys(
                 validated_X, search_keys, is_demo_dataset, is_transform=True, silent_mode=silent_mode
             )
@@ -2153,7 +2171,7 @@ class FeaturesEnricher(TransformerMixin):
             df = self.__add_country_code(df, search_keys)
 
             generated_features = []
-            date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+            date_column = self._get_date_column(search_keys)
             if date_column is not None:
                 converter = DateTimeSearchKeyConverter(date_column, self.date_format, self.logger, bundle=self.bundle)
                 df = converter.convert(df, keep_time=True)
@@ -2161,7 +2179,7 @@ class FeaturesEnricher(TransformerMixin):
                 generated_features.extend(converter.generated_features)
             else:
                 self.logger.info("Input dataset hasn't date column")
-                if self.add_date_if_missing:
+                if self.__should_add_date_column():
                     df = self._add_current_date_as_key(df, search_keys, self.logger, self.bundle)
 
             email_columns = SearchKey.find_all_keys(search_keys, SearchKey.EMAIL)
@@ -2292,6 +2310,7 @@ class FeaturesEnricher(TransformerMixin):
                 meaning_types=meaning_types,
                 search_keys=combined_search_keys,
                 unnest_search_keys=unnest_search_keys,
+                id_columns=self.__get_renamed_id_columns(columns_renaming),
                 date_format=self.date_format,
                 rest_client=self.rest_client,
                 logger=self.logger,
@@ -2444,7 +2463,14 @@ class FeaturesEnricher(TransformerMixin):
         # 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]:
+            if multi_key not in [
+                SearchKey.PHONE,
+                SearchKey.IP,
+                SearchKey.POSTAL_CODE,
+                SearchKey.EMAIL,
+                SearchKey.HEM,
+                SearchKey.CUSTOM_KEY,
+            ]:
                 msg = self.bundle.get("unsupported_multi_key").format(multi_key)
                 self.logger.warning(msg)
                 raise ValidationError(msg)
@@ -2608,7 +2634,7 @@ class FeaturesEnricher(TransformerMixin):
             self.fit_generated_features.extend(converter.generated_features)
         else:
             self.logger.info("Input dataset hasn't date column")
-            if self.add_date_if_missing:
+            if self.__should_add_date_column():
                 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)
@@ -2641,6 +2667,13 @@ class FeaturesEnricher(TransformerMixin):
 
         self.__adjust_cv(df)
 
+        if self.id_columns is not None and self.cv is not None and self.cv.is_time_series():
+            id_columns = self.__get_renamed_id_columns()
+            if id_columns:
+                self.fit_search_keys.update({col: SearchKey.CUSTOM_KEY for col in id_columns})
+                self.search_keys.update({col: SearchKey.CUSTOM_KEY for col in self.id_columns})
+                self.runtime_parameters.properties["id_columns"] = ",".join(id_columns)
+
         df, fintech_warnings = remove_fintech_duplicates(
             df, self.fit_search_keys, date_format=self.date_format, logger=self.logger, bundle=self.bundle
         )
@@ -2762,6 +2795,8 @@ class FeaturesEnricher(TransformerMixin):
             search_keys=combined_search_keys,
             unnest_search_keys=unnest_search_keys,
             model_task_type=self.model_task_type,
+            cv_type=self.cv,
+            id_columns=self.__get_renamed_id_columns(),
             date_format=self.date_format,
             random_state=self.random_state,
             rest_client=self.rest_client,
@@ -2918,6 +2953,14 @@ class FeaturesEnricher(TransformerMixin):
         if not self.warning_counter.has_warnings():
             self.__display_support_link(self.bundle.get("all_ok_community_invite"))
 
+    def __should_add_date_column(self):
+        return self.add_date_if_missing or (self.cv is not None and self.cv.is_time_series())
+
+    def __get_renamed_id_columns(self, renaming: Optional[Dict[str, str]] = None):
+        renaming = renaming or self.fit_columns_renaming
+        reverse_renaming = {v: k for k, v in renaming.items()}
+        return None if self.id_columns is None else [reverse_renaming.get(c) or c for c in self.id_columns]
+
     def __adjust_cv(self, df: pd.DataFrame):
         date_column = SearchKey.find_key(self.fit_search_keys, [SearchKey.DATE, SearchKey.DATETIME])
         # Check Multivariate time series
@@ -3163,7 +3206,7 @@ class FeaturesEnricher(TransformerMixin):
             if DateTimeSearchKeyConverter.DATETIME_COL in X.columns:
                 date_column = DateTimeSearchKeyConverter.DATETIME_COL
             else:
-                date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+                date_column = FeaturesEnricher._get_date_column(search_keys)
             sort_columns = [date_column] if date_column is not None else []
 
             # Xy = pd.concat([X, y], axis=1)
@@ -3226,6 +3269,7 @@ class FeaturesEnricher(TransformerMixin):
                 f"Generate features: {self.generate_features}\n"
                 f"Round embeddings: {self.round_embeddings}\n"
                 f"Detect missing search keys: {self.detect_missing_search_keys}\n"
+                f"Exclude columns: {self.exclude_columns}\n"
                 f"Exclude features sources: {exclude_features_sources}\n"
                 f"Calculate metrics: {calculate_metrics}\n"
                 f"Scoring: {scoring}\n"
@@ -3233,6 +3277,15 @@ class FeaturesEnricher(TransformerMixin):
                 f"Remove target outliers: {remove_outliers_calc_metrics}\n"
                 f"Exclude columns: {self.exclude_columns}\n"
                 f"Search id: {self.search_id}\n"
+                f"Custom loss: {self.loss}\n"
+                f"Logs enabled: {self.logs_enabled}\n"
+                f"Raise validation error: {self.raise_validation_error}\n"
+                f"Baseline score column: {self.baseline_score_column}\n"
+                f"Client ip: {self.client_ip}\n"
+                f"Client visitorId: {self.client_visitorid}\n"
+                f"Add date if missing: {self.add_date_if_missing}\n"
+                f"Select features: {self.select_features}\n"
+                f"Disable force downsampling: {self.disable_force_downsampling}\n"
             )
 
             def sample(df):
@@ -3355,6 +3408,10 @@ class FeaturesEnricher(TransformerMixin):
             if t == SearchKey.POSTAL_CODE:
                 return col
 
+    @staticmethod
+    def _get_date_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+        return SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+
     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]]]:
@@ -3363,7 +3420,9 @@ class FeaturesEnricher(TransformerMixin):
         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
+            key_type: key_names
+            for key_type, key_names in search_key_names_by_type.items()
+            if len(key_names) > 1 and key_type != SearchKey.CUSTOM_KEY
         }
         if len(search_key_names_by_type) == 0:
             return df, {}
@@ -3416,9 +3475,9 @@ class FeaturesEnricher(TransformerMixin):
             ]
             if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
                 date_column = DateTimeSearchKeyConverter.DATETIME_COL
-                sort_exclude_columns.append(SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME]))
+                sort_exclude_columns.append(self._get_date_column(search_keys))
             else:
-                date_column = SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
+                date_column = self._get_date_column(search_keys)
             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)))
@@ -3854,11 +3913,6 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(msg + f" Provided search keys: {search_keys}")
             raise ValidationError(msg)
 
-        if SearchKey.CUSTOM_KEY in valid_search_keys.values():
-            custom_keys = [column for column, key in valid_search_keys.items() if key == SearchKey.CUSTOM_KEY]
-            for key in custom_keys:
-                del valid_search_keys[key]
-
         if (
             len(valid_search_keys.values()) == 1
             and self.country_code is None

{upgini-1.2.36a1 → upgini-1.2.38}/src/upgini/metadata.py

@@ -350,3 +350,6 @@ class CVType(Enum):
     time_series = "time_series"
     blocked_time_series = "blocked_time_series"
     not_set = "not_set"
+
+    def is_time_series(self) -> bool:
+        return self in [CVType.time_series, CVType.blocked_time_series]

{upgini-1.2.36a1 → upgini-1.2.38}/src/upgini/resource_bundle/strings.properties

@@ -216,6 +216,7 @@ imbalanced_target=\nTarget is imbalanced and will be undersampled. Frequency of
 loss_selection_info=Using loss `{}` for feature selection
 loss_calc_metrics_info=Using loss `{}` for metrics calculation with default estimator
 forced_balance_undersample=For quick data retrieval, your dataset has been sampled. To use data search without data sampling please contact support (sales@upgini.com)
+online_api_features_transform=Please note that some of the selected features {} are provided through a slow enrichment interface and are not available via transformation. However, they can be accessed via the API:
 
 # Validation table
 validation_column_name_header=Column name

{upgini-1.2.36a1 → upgini-1.2.38}/src/upgini/utils/target_utils.py

@@ -1,15 +1,18 @@
+import itertools
 import logging
-from typing import Callable, Optional, Union
+from typing import Callable, List, Optional, Union
 
 import numpy as np
 import pandas as pd
 from pandas.api.types import is_numeric_dtype, is_bool_dtype
 
 from upgini.errors import ValidationError
-from upgini.metadata import SYSTEM_RECORD_ID, ModelTaskType
+from upgini.metadata import SYSTEM_RECORD_ID, CVType, ModelTaskType
 from upgini.resource_bundle import ResourceBundle, bundle, get_custom_bundle
 from upgini.sampler.random_under_sampler import RandomUnderSampler
 
+TS_MIN_DIFFERENT_IDS_RATIO = 0.2
+
 
 def correct_string_target(y: Union[pd.Series, np.ndarray]) -> Union[pd.Series, np.ndarray]:
     if isinstance(y, pd.Series):
@@ -201,7 +204,10 @@ def balance_undersample(
 def balance_undersample_forced(
     df: pd.DataFrame,
     target_column: str,
+    id_columns: List[str],
+    date_column: str,
     task_type: ModelTaskType,
+    cv_type: CVType | None,
     random_state: int,
     sample_size: int = 7000,
     logger: Optional[logging.Logger] = None,
@@ -233,7 +239,17 @@ def balance_undersample_forced(
 
     resampled_data = df
     df = df.copy().sort_values(by=SYSTEM_RECORD_ID)
-    if task_type in [ModelTaskType.MULTICLASS, ModelTaskType.REGRESSION, ModelTaskType.TIMESERIES]:
+    if cv_type is not None and cv_type.is_time_series():
+        logger.warning(f"Sampling time series dataset from {len(df)} to {sample_size}")
+        resampled_data = balance_undersample_time_series(
+            df,
+            id_columns=id_columns,
+            date_column=date_column,
+            sample_size=sample_size,
+            random_state=random_state,
+            logger=logger,
+        )
+    elif task_type in [ModelTaskType.MULTICLASS, ModelTaskType.REGRESSION]:
         logger.warning(f"Sampling dataset from {len(df)} to {sample_size}")
         resampled_data = df.sample(n=sample_size, random_state=random_state)
     else:
@@ -264,6 +280,65 @@ def balance_undersample_forced(
     return resampled_data
 
 
+def balance_undersample_time_series(
+    df: pd.DataFrame,
+    id_columns: List[str],
+    date_column: str,
+    sample_size: int,
+    random_state: int = 42,
+    min_different_ids_ratio: float = TS_MIN_DIFFERENT_IDS_RATIO,
+    prefer_recent_dates: bool = True,
+    logger: Optional[logging.Logger] = None,
+):
+    def ensure_tuple(x):
+        return tuple([x]) if not isinstance(x, tuple) else x
+
+    random_state = np.random.RandomState(random_state)
+
+    if not id_columns:
+        id_columns = [date_column]
+    ids_sort = df.groupby(id_columns)[date_column].aggregate(["max", "count"]).T.to_dict()
+    ids_sort = {
+        ensure_tuple(k): (
+            (v["max"], v["count"], random_state.rand()) if prefer_recent_dates else (v["count"], random_state.rand())
+        )
+        for k, v in ids_sort.items()
+    }
+    id_counts = df[id_columns].value_counts()
+    id_counts.index = [ensure_tuple(i) for i in id_counts.index]
+    id_counts = id_counts.sort_index(key=lambda x: [ids_sort[y] for y in x], ascending=False).cumsum()
+    id_counts = id_counts[id_counts <= sample_size]
+    min_different_ids = max(int(len(df[id_columns].drop_duplicates()) * min_different_ids_ratio), 1)
+
+    def id_mask(sample_index: pd.Index) -> pd.Index:
+        if isinstance(sample_index, pd.MultiIndex):
+            return pd.MultiIndex.from_frame(df[id_columns]).isin(sample_index)
+        else:
+            return df[id_columns[0]].isin(sample_index)
+
+    if len(id_counts) < min_different_ids:
+        if logger is not None:
+            logger.info(
+                f"Different ids count {len(id_counts)} for sample size {sample_size} is less than min different ids {min_different_ids}, sampling time window"
+            )
+        date_counts = df.groupby(id_columns)[date_column].nunique().sort_values(ascending=False)
+        ids_to_sample = date_counts.index[:min_different_ids] if len(id_counts) > 0 else date_counts.index
+        mask = id_mask(ids_to_sample)
+        df = df[mask]
+        sample_date_counts = df[date_column].value_counts().sort_index(ascending=False).cumsum()
+        sample_date_counts = sample_date_counts[sample_date_counts <= sample_size]
+        df = df[df[date_column].isin(sample_date_counts.index)]
+    else:
+        if len(id_columns) > 1:
+            id_counts.index = pd.MultiIndex.from_tuples(id_counts.index)
+        else:
+            id_counts.index = [i[0] for i in id_counts.index]
+        mask = id_mask(id_counts.index)
+        df = df[mask]
+
+    return df
+
+
 def calculate_psi(expected: pd.Series, actual: pd.Series) -> Union[float, Exception]:
     try:
         df = pd.concat([expected, actual])

upgini-1.2.36a1/src/upgini/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.2.36a1"