upgini 1.2.157.dev2__tar.gz → 1.2.157.dev3__tar.gz

{upgini-1.2.157.dev2 → upgini-1.2.157.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.157.dev2
+Version: 1.2.157.dev3
 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.157.dev3/src/upgini/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.2.157.dev3"

{upgini-1.2.157.dev2 → upgini-1.2.157.dev3}/src/upgini/autofe/feature.py

@@ -8,7 +8,7 @@ import pandas as pd
 from pandas._typing import DtypeObj
 
 from upgini.autofe.all_operators import find_op
-from upgini.autofe.operand import CalculationContext, OperandValue
+from upgini.autofe.operand import CalculationContext, OperandValue, operand_to_output
 from upgini.autofe.operator import Operator, PandasOperator
 from upgini.autofe.timeseries.base import TimeSeriesBase
 from upgini.autofe.utils import pydantic_dump_method, pydantic_parse_method
@@ -62,9 +62,14 @@ class Column:
     def infer_type(self, data: pd.DataFrame) -> DtypeObj:
         return data[self.name].dtype
 
-    def calculate(self, data: Union[pd.DataFrame, CalculationContext]) -> pd.Series:
+    def calculate(
+        self,
+        data: Union[pd.DataFrame, CalculationContext],
+        preserve_kind: bool = False,
+    ) -> Union[pd.Series, OperandValue]:
         ctx = data if isinstance(data, CalculationContext) else CalculationContext.from_dataframe(data)
-        return self._eval(ctx).as_series()
+        operand = self._eval(ctx)
+        return operand if preserve_kind else operand.as_series()
 
     def _eval(self, ctx: CalculationContext) -> OperandValue:
         self.data = ctx.resolve_operand(self.name)
@@ -245,10 +250,11 @@ class Feature:
         self,
         data: Union[pd.DataFrame, CalculationContext],
         is_root: bool = False,
-    ) -> Union[pd.Series, pd.DataFrame]:
+        preserve_kind: bool = False,
+    ) -> Union[pd.Series, pd.DataFrame, OperandValue]:
         ctx = data if isinstance(data, CalculationContext) else CalculationContext.from_dataframe(data)
         result = self._eval(ctx)
-        new_data = self._operand_to_output(result)
+        new_data = operand_to_output(result, preserve_kind=preserve_kind)
 
         if is_root:
             self.data = new_data
@@ -266,13 +272,6 @@ class Feature:
             return self.op.calculate(left=left, right=right)
         raise NotImplementedError(f"Unrecognized operator {self.op.name}.")
 
-    @staticmethod
-    def _operand_to_output(operand: OperandValue) -> Union[pd.Series, pd.DataFrame]:
-        new_data = operand.as_series()
-        if (str(new_data.dtype) == "category") | (str(new_data.dtype) == "object"):
-            return new_data
-        return new_data.replace([-np.inf, np.inf], np.nan)
-
     @staticmethod
     def check_xor(left: Union[Column, "Feature"], right: Union[Column, "Feature"]) -> bool:
         def _get_all_columns(feature: Union[Column, "Feature"]) -> List[str]:

{upgini-1.2.157.dev2 → upgini-1.2.157.dev3}/src/upgini/autofe/operand.py

@@ -94,6 +94,27 @@ class CalculationContext:
         raise KeyError(name)
 
 
+def finalize_operand(operand: OperandValue) -> OperandValue:
+    if operand.kind == OperandKind.MATRIX:
+        matrix = np.where(np.isinf(operand.as_matrix()), np.nan, operand.as_matrix())
+        return OperandValue.from_matrix(matrix, index=operand.index, source=operand.source)
+    if operand.kind == OperandKind.ARRAY:
+        array = operand.data  # type: ignore[assignment]
+        array = np.where(np.isinf(array), np.nan, array)
+        return OperandValue.from_array(array, index=operand.index, source=operand.source)
+    new_data = operand.as_series()
+    if (str(new_data.dtype) == "category") | (str(new_data.dtype) == "object"):
+        return OperandValue.from_series(new_data, source=operand.source)
+    return OperandValue.from_series(new_data.replace([-np.inf, np.inf], np.nan), source=operand.source)
+
+
+def operand_to_output(operand: OperandValue, preserve_kind: bool = False) -> Union[pd.Series, OperandValue]:
+    finalized = finalize_operand(operand)
+    if preserve_kind:
+        return finalized
+    return finalized.as_series()
+
+
 def wrap_operand(
     result: Union[pd.Series, np.ndarray, "OperandValue"],
     source: Optional[str] = None,

{upgini-1.2.157.dev2 → upgini-1.2.157.dev3}/src/upgini/autofe/unary.py

@@ -162,7 +162,7 @@ class Embeddings(PandasOperator):
     name: str = "emb"
     is_unary: bool = True
     input_type: Optional[str] = "string"
-    output_type: Optional[str] = "matrix"
+    output_type: Optional[str] = "vector"
 
 
 class Bin(PandasOperator):

upgini-1.2.157.dev2/src/upgini/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.2.157.dev2"