upgini 1.1.275__py3-none-any.whl → 1.1.275a1__py3-none-any.whl

upgini/ads.py

@@ -5,7 +5,7 @@ from typing import Dict, Optional
 
 import numpy as np
 import pandas as pd
-from pandas.api.types import is_object_dtype, is_string_dtype
+from pandas.api.types import is_string_dtype
 
 from upgini import SearchKey
 from upgini.http import get_rest_client
@@ -34,11 +34,7 @@ def upload_user_ads(name: str, df: pd.DataFrame, search_keys: Dict[str, SearchKe
             if df[column_name].notnull().sum() < min_valid_rows_count:
                 raise ValueError(bundle.get("ads_upload_to_many_empty_rows"))
             meaning_type = search_keys[column_name].value
-            if (
-                meaning_type == FileColumnMeaningType.MSISDN
-                and not is_string_dtype(df[column_name])
-                and not is_object_dtype(df[column_name])
-            ):
+            if meaning_type == FileColumnMeaningType.MSISDN and not is_string_dtype(df[column_name]):
                 df[column_name] = df[column_name].values.astype(np.int64).astype("string")  # type: ignore
         else:
             meaning_type = FileColumnMeaningType.FEATURE

upgini/autofe/date.py

@@ -2,7 +2,6 @@ from typing import Any, Optional, Union
 import numpy as np
 import pandas as pd
 from pydantic import BaseModel
-from pandas.core.arrays.timedeltas import TimedeltaArray
 
 from upgini.autofe.operand import PandasOperand
 
@@ -47,7 +46,6 @@ class DateDiffType2(PandasOperand, DateDiffMixin):
         future = right + (left.dt.year - right.dt.year).apply(
             lambda y: np.datetime64("NaT") if np.isnan(y) else pd.tseries.offsets.DateOffset(years=y)
         )
-        future = pd.to_datetime(future)
         before = future[future < left]
         future[future < left] = before + pd.tseries.offsets.DateOffset(years=1)
         diff = (future - left) / np.timedelta64(1, self.diff_unit)
@@ -74,13 +72,8 @@ class DateListDiff(PandasOperand, DateDiffMixin):
 
         return pd.Series(left - right.values).apply(lambda x: self._agg(self._diff(x)))
 
-    def _diff(self, x: TimedeltaArray):
-        if self.diff_unit == "Y":
-            x = (x / 365 / 24 / 60 / 60 / 10**9).astype(int)
-        elif self.diff_unit == "M":
-            raise Exception("Unsupported difference unit: Month")
-        else:
-            x = x / np.timedelta64(1, self.diff_unit)
+    def _diff(self, x):
+        x = x / np.timedelta64(1, self.diff_unit)
         return x[x > 0]
 
     def _agg(self, x):

upgini/data_source/data_source_publisher.py

@@ -48,7 +48,6 @@ class DataSourcePublisher:
         data_table_uri: str,
         search_keys: Dict[str, SearchKey],
         update_frequency: str,
-        exclude_from_autofe_generation: Optional[List[str]],
         secondary_search_keys: Optional[Dict[str, SearchKey]] = None,
         sort_column: Optional[str] = None,
         date_format: Optional[str] = None,
@@ -58,6 +57,7 @@ class DataSourcePublisher:
         join_date_abs_limit_days: Optional[int] = None,
         features_for_embeddings: Optional[List[str]] = DEFAULT_GENERATE_EMBEDDINGS,
         data_table_id_to_replace: Optional[str] = None,
+        exclude_from_autofe_generation: Optional[List[str]] = None,
         _force_generation=False,
         _silent=False,
     ) -> str:

upgini/dataset.py

@@ -17,13 +17,14 @@ from pandas.api.types import (
     is_numeric_dtype,
     is_period_dtype,
     is_string_dtype,
-    is_object_dtype,
 )
 
 from upgini.errors import ValidationError
 from upgini.http import ProgressStage, SearchProgress, _RestClient
 from upgini.metadata import (
+    ENTITY_SYSTEM_RECORD_ID,
     EVAL_SET_INDEX,
+    SEARCH_KEY_UNNEST,
     SYSTEM_COLUMNS,
     SYSTEM_RECORD_ID,
     TARGET,
@@ -79,6 +80,7 @@ class Dataset:  # (pd.DataFrame):
         path: Optional[str] = None,
         meaning_types: Optional[Dict[str, FileColumnMeaningType]] = None,
         search_keys: Optional[List[Tuple[str, ...]]] = None,
+        unnest_search_keys: Optional[List[str]] = None,
         model_task_type: Optional[ModelTaskType] = None,
         random_state: Optional[int] = None,
         rest_client: Optional[_RestClient] = None,
@@ -113,6 +115,7 @@ class Dataset:  # (pd.DataFrame):
         self.description = description
         self.meaning_types = meaning_types
         self.search_keys = search_keys
+        self.unnest_search_keys = unnest_search_keys
         self.ignore_columns = []
         self.hierarchical_group_keys = []
         self.hierarchical_subgroup_keys = []
@@ -172,7 +175,7 @@ class Dataset:  # (pd.DataFrame):
         new_columns = []
         dup_counter = 0
         for column in self.data.columns:
-            if column in [TARGET, EVAL_SET_INDEX, SYSTEM_RECORD_ID]:
+            if column in [TARGET, EVAL_SET_INDEX, SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID, SEARCH_KEY_UNNEST]:
                 self.columns_renaming[column] = column
                 new_columns.append(column)
                 continue
@@ -220,7 +223,7 @@ class Dataset:  # (pd.DataFrame):
         """Check that string values less than maximum characters for LLM"""
         # self.logger.info("Validate too long string values")
         for col in self.data.columns:
-            if is_string_dtype(self.data[col]) or is_object_dtype(self.data[col]):
+            if is_string_dtype(self.data[col]):
                 max_length: int = self.data[col].astype("str").str.len().max()
                 if max_length > self.MAX_STRING_FEATURE_LENGTH:
                     self.data[col] = self.data[col].astype("str").str.slice(stop=self.MAX_STRING_FEATURE_LENGTH)
@@ -351,9 +354,11 @@ class Dataset:  # (pd.DataFrame):
         if postal_code is not None and postal_code in self.data.columns:
             # self.logger.info("Normalize postal code")
 
-            if is_string_dtype(self.data[postal_code]) or is_object_dtype(self.data[postal_code]):
+            if is_string_dtype(self.data[postal_code]):
                 try:
-                    self.data[postal_code] = self.data[postal_code].astype("float64").astype("Int64").astype("string")
+                    self.data[postal_code] = (
+                        self.data[postal_code].astype("string").astype("Float64").astype("Int64").astype("string")
+                    )
                 except Exception:
                     pass
             elif is_float_dtype(self.data[postal_code]):
@@ -803,6 +808,8 @@ class Dataset:  # (pd.DataFrame):
                     meaningType=meaning_type,
                     minMaxValues=min_max_values,
                 )
+                if self.unnest_search_keys and column_meta.originalName in self.unnest_search_keys:
+                    column_meta.isUnnest = True
 
                 columns.append(column_meta)
 
@@ -822,7 +829,7 @@ class Dataset:  # (pd.DataFrame):
             return DataType.INT
         elif is_float_dtype(pandas_data_type):
             return DataType.DECIMAL
-        elif is_string_dtype(pandas_data_type) or is_object_dtype(pandas_data_type):
+        elif is_string_dtype(pandas_data_type):
             return DataType.STRING
         else:
             msg = self.bundle.get("dataset_invalid_column_type").format(column_name, pandas_data_type)