upgini 1.2.156.dev5__tar.gz → 1.2.157a1.dev1__tar.gz

{upgini-1.2.156.dev5 → upgini-1.2.157a1.dev1}/PKG-INFO

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

@@ -0,0 +1 @@
+__version__ = "1.2.157.a1.dev1"

{upgini-1.2.156.dev5 → upgini-1.2.157a1.dev1}/src/upgini/autofe/binary.py

@@ -140,120 +140,23 @@ class Distance(PandasOperator):
     is_symmetrical: bool = True
     has_symmetry_importance: bool = True
 
-    def _to_float_vector(self, value) -> Optional[np.ndarray]:
-        if value is None or (isinstance(value, float) and np.isnan(value)):
-            return None
-        return np.atleast_1d(np.asarray(value, dtype=np.float64))
-
-    def _stack_vectors(self, arr: np.ndarray, indices: np.ndarray) -> Optional[np.ndarray]:
-        try:
-            stacked = np.vstack(arr[indices])
-        except (ValueError, TypeError):
-            return None
-        if stacked.ndim != 2 or stacked.shape[1] == 0:
-            return None
-        if stacked.dtype != np.float64:
-            stacked = stacked.astype(np.float64, copy=False)
-        return stacked
-
-    def _cosine_distance_stacked(self, left_mat: np.ndarray, right_mat: np.ndarray) -> np.ndarray:
-        if left_mat.shape != right_mat.shape:
-            raise ValueError("Left and right matrices must have the same shape")
-        dots = np.sum(left_mat * right_mat, axis=1)
-        denom = np.sqrt(np.sum(left_mat * left_mat, axis=1)) * np.sqrt(np.sum(right_mat * right_mat, axis=1))
-        results = np.full(len(left_mat), np.nan, dtype=np.float64)
-        valid = denom > 0
-        results[valid] = 1 - dots[valid] / denom[valid]
-        return results
-
-    def _try_fast_distance(
-        self, left_arr: np.ndarray, right_arr: np.ndarray, row_present: np.ndarray, n: int
-    ) -> Optional[np.ndarray]:
-        valid_idx = np.flatnonzero(row_present)
-        if len(valid_idx) == 0:
-            return None
-
-        left_mat = self._stack_vectors(left_arr, valid_idx)
-        right_mat = self._stack_vectors(right_arr, valid_idx)
-        if left_mat is None or right_mat is None:
-            return None
-
-        stacked_results = self._cosine_distance_stacked(left_mat, right_mat)
-        if row_present.all():
-            return stacked_results
-
-        results = np.full(n, np.nan, dtype=np.float64)
-        results[valid_idx] = stacked_results
-        return results
-
-    def _parse_vectors(self, left_arr: np.ndarray, right_arr: np.ndarray) -> tuple[list, list, np.ndarray]:
-        n = len(left_arr)
-        left_vecs: list[Optional[np.ndarray]] = [None] * n
-        right_vecs: list[Optional[np.ndarray]] = [None] * n
-        row_valid = np.zeros(n, dtype=bool)
-        for i, (left_value, right_value) in enumerate(zip(left_arr, right_arr)):
-            left_vec = self._to_float_vector(left_value)
-            right_vec = self._to_float_vector(right_value)
-            left_vecs[i] = left_vec
-            right_vecs[i] = right_vec
-            if left_vec is not None and right_vec is not None and left_vec.size > 0 and right_vec.size > 0:
-                row_valid[i] = True
-        return left_vecs, right_vecs, row_valid
-
-    def _masked_cosine_distance(
-        self, left_mat: np.ndarray, right_mat: np.ndarray, row_valid: np.ndarray
-    ) -> np.ndarray:
-        results = self._cosine_distance_stacked(left_mat, right_mat)
-        if len(results) == len(row_valid):
-            results[~row_valid] = np.nan
-        return results
-
-    def _calculate_rowwise(
-        self, results: np.ndarray, indices: np.ndarray, left_vecs: list, right_vecs: list
-    ) -> None:
-        for index in indices:
-            left_vec = left_vecs[index]
-            right_vec = right_vecs[index]
-            dot = np.sum(left_vec * right_vec)
-            denom = np.sqrt(np.sum(left_vec * left_vec)) * np.sqrt(np.sum(right_vec * right_vec))
-            if denom != 0:
-                results[index] = 1 - dot / denom
-
-    def _calculate_fallback_distance(
-        self, left_arr: np.ndarray, right_arr: np.ndarray, n: int, index: pd.Index
-    ) -> pd.Series:
-        results = np.full(n, np.nan, dtype=np.float64)
-        left_vecs, right_vecs, row_valid = self._parse_vectors(left_arr, right_arr)
-        if not row_valid.any():
-            return pd.Series(results, index=index, dtype=np.float64)
-
-        valid_idx = np.flatnonzero(row_valid)
-        dim = left_vecs[valid_idx[0]].size
-        if not all(left_vecs[i].size == dim and right_vecs[i].size == dim for i in valid_idx):
-            self._calculate_rowwise(results, valid_idx, left_vecs, right_vecs)
-            return pd.Series(results, index=index, dtype=np.float64)
-
-        left_mat = np.zeros((n, dim), dtype=np.float64)
-        right_mat = np.zeros((n, dim), dtype=np.float64)
-        left_mat[valid_idx] = np.stack([left_vecs[i] for i in valid_idx])
-        right_mat[valid_idx] = np.stack([right_vecs[i] for i in valid_idx])
-
-        return pd.Series(self._masked_cosine_distance(left_mat, right_mat, row_valid), index=index, dtype=np.float64)
-
     def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
-        n = len(left)
-        if n == 0:
-            return pd.Series(np.full(n, np.nan), index=left.index, dtype=np.float64)
-
-        row_present = left.notna().to_numpy() & right.notna().to_numpy()
-        left_arr = left.to_numpy()
-        right_arr = right.to_numpy()
-
-        fast_results = self._try_fast_distance(left_arr, right_arr, row_present, n)
-        if fast_results is not None:
-            return pd.Series(fast_results, index=left.index, dtype=np.float64)
-
-        return self._calculate_fallback_distance(left_arr, right_arr, n, left.index)
+        return pd.Series(
+            1 - self.__dot(left, right) / (self.__norm(left) * self.__norm(right)), index=left.index
+        ).astype(np.float64)
+
+    # row-wise dot product, handling None values
+    def __dot(self, left: pd.Series, right: pd.Series) -> pd.Series:
+        left = left.apply(lambda x: np.array(x).astype(np.float64))
+        right = right.apply(lambda x: np.array(x).astype(np.float64))
+        res = (left.dropna() * right.dropna()).apply(np.sum)
+        res = res.reindex(left.index.union(right.index))
+        return res
+
+    # Calculate the norm of a vector, handling None values
+    def __norm(self, vector: pd.Series) -> pd.Series:
+        vector = vector.fillna(np.nan)
+        return np.sqrt(self.__dot(vector, vector))
 
 
 # Left for backward compatibility

{upgini-1.2.156.dev5 → upgini-1.2.157a1.dev1}/src/upgini/autofe/date.py

@@ -8,7 +8,7 @@ from pandas.core.arrays.timedeltas import TimedeltaArray
 from pydantic import BaseModel, __version__ as pydantic_version
 
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator
-from upgini.autofe.utils import bin_index, bin_index_many, bin_index_vectorized, pydantic_validator
+from upgini.autofe.utils import pydantic_validator
 
 
 def get_pydantic_version():
@@ -113,172 +113,6 @@ _ext_aggregations = {"nunique": (lambda x: len(np.unique(x)), 0), "count": (len,
 _count_aggregations = ["nunique", "count"]
 
 
-def _aggregate_diffs(values: np.ndarray, aggregation: str) -> float:
-    values = np.atleast_1d(np.asarray(values, dtype=np.float64))
-    method = getattr(np, aggregation, None)
-    default = np.nan
-    if method is None and aggregation in _ext_aggregations:
-        method, default = _ext_aggregations[aggregation]
-    elif not callable(method):
-        raise ValueError(f"Unsupported aggregation: {aggregation}")
-
-    return method(values) if len(values) > 0 else default
-
-
-class DateListDiffLists(PandasOperator, DateDiffMixin, ParametrizedOperator):
-    name: str = "date_diff_lists"
-    is_binary: bool = True
-    has_symmetry_importance: bool = True
-    output_type: Optional[str] = "vector"
-
-    replace_negative: bool = False
-
-    def get_params(self) -> Dict[str, Optional[str]]:
-        res = super().get_params()
-        res.update(
-            {
-                "diff_unit": self.diff_unit,
-                "left_unit": self.left_unit,
-                "right_unit": self.right_unit,
-                "replace_negative": self.replace_negative,
-            }
-        )
-        return res
-
-    def to_formula(self) -> str:
-        return f"date_diff_lists_{self.diff_unit}"
-
-    @classmethod
-    def from_formula(cls, formula: str) -> Optional["DateListDiffLists"]:
-        if formula == "date_diff_lists":
-            return cls()
-        if formula.startswith("date_diff_lists_"):
-            diff_unit = formula.replace("date_diff_lists_", "")
-            if diff_unit in {"D", "Y"}:
-                return cls(diff_unit=diff_unit)
-        return None
-
-    @staticmethod
-    def _non_empty_list_mask(right: pd.Series) -> pd.Series:
-        values = right.to_numpy()
-        mask = np.empty(len(values), dtype=bool)
-        for i, value in enumerate(values):
-            if value is None or (isinstance(value, float) and np.isnan(value)):
-                mask[i] = False
-            elif isinstance(value, (list, tuple, np.ndarray)):
-                mask[i] = len(value) > 0
-            else:
-                mask[i] = False
-        return pd.Series(mask, index=right.index)
-
-    def _convert_date_lists(self, lists: pd.Series) -> pd.Series:
-        exploded = lists.explode()
-        converted = pd.to_datetime(exploded, unit=self.right_unit, errors="coerce")
-        return pd.Series(
-            {
-                idx: pd.arrays.DatetimeArray(values.to_numpy())
-                for idx, values in converted.groupby(converted.index, sort=False)
-            }
-        )
-
-    def _row_diffs(self, left_date, right_dates: pd.arrays.DatetimeArray) -> List[float]:
-        diffs = self._convert_diff_to_unit(left_date - right_dates)
-        if self.replace_negative:
-            diffs = diffs[diffs > 0]
-        return np.atleast_1d(np.asarray(diffs, dtype=np.float64)).tolist()
-
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
-        if left.isna().all() or right.isna().all():
-            return pd.Series([None] * len(left), index=left.index, dtype=object)
-
-        left = self._convert_to_date(left, self.left_unit)
-        right_mask = self._non_empty_list_mask(right)
-        mask = left.notna() & right.notna() & right_mask
-
-        results = pd.Series([None] * len(left), index=left.index, dtype=object)
-        if not mask.any():
-            return results
-
-        masked_left = left[mask]
-        converted_lists = self._convert_date_lists(right[mask])
-        for idx, left_date in masked_left.items():
-            results.loc[idx] = self._row_diffs(left_date, converted_lists[idx])
-        return results
-
-
-class DateListDiffAggWithinBounds(PandasOperator, ParametrizedOperator):
-    name: str = "date_diff_list_agg"
-    is_unary: bool = True
-    output_type: Optional[str] = "float"
-
-    lower_bound: Optional[int] = None
-    upper_bound: Optional[int] = None
-    aggregation: str
-    normalize: Optional[bool] = None
-
-    def get_params(self) -> Dict[str, Optional[str]]:
-        res = super().get_params()
-        res.update(
-            {
-                "aggregation": self.aggregation,
-                "lower_bound": str(self.lower_bound) if self.lower_bound is not None else None,
-                "upper_bound": str(self.upper_bound) if self.upper_bound is not None else None,
-                "normalize": str(self.normalize) if self.normalize is not None else None,
-            }
-        )
-        return res
-
-    def to_formula(self) -> str:
-        lower_bound = "minusinf" if self.lower_bound is None else self.lower_bound
-        upper_bound = "plusinf" if self.upper_bound is None else self.upper_bound
-        norm = "_norm" if self.normalize else ""
-        return f"date_diff_list_agg_{lower_bound}_{upper_bound}_{self.aggregation}{norm}"
-
-    @classmethod
-    def from_formula(cls, formula: str) -> Optional["DateListDiffAggWithinBounds"]:
-        import re
-
-        normalize = formula.endswith("_norm")
-        formula = formula.replace("_norm", "")
-
-        pattern = r"^date_diff_list_agg_((minusinf|\d+))_((plusinf|\d+))_(\w+)$"
-        match = re.match(pattern, formula)
-        if not match:
-            return None
-
-        lower_bound = None if match.group(1) == "minusinf" else int(match.group(1))
-        upper_bound = None if match.group(3) == "plusinf" else int(match.group(3))
-        aggregation = match.group(5)
-        return cls(
-            lower_bound=lower_bound,
-            upper_bound=upper_bound,
-            aggregation=aggregation,
-            normalize=normalize,
-        )
-
-    def _aggregate_row(self, diffs) -> float:
-        if diffs is None or (isinstance(diffs, float) and np.isnan(diffs)):
-            return np.nan
-
-        diffs = np.atleast_1d(np.asarray(diffs, dtype=np.float64))
-        orig_len = len(diffs)
-        if self.lower_bound is not None or self.upper_bound is not None:
-            lower = self.lower_bound if self.lower_bound is not None else -np.inf
-            upper = self.upper_bound if self.upper_bound is not None else np.inf
-            diffs = diffs[(diffs >= lower) & (diffs < upper)]
-        agg_res = _aggregate_diffs(diffs, self.aggregation)
-        if self.normalize and orig_len > 0:
-            return agg_res / orig_len
-        return agg_res
-
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
-        results = np.empty(len(data), dtype=np.float64)
-        results[:] = np.nan
-        for i, diffs in enumerate(data.to_numpy()):
-            results[i] = self._aggregate_row(diffs)
-        return pd.Series(results, index=data.index, dtype=np.float64)
-
-
 class DateListDiff(PandasOperator, DateDiffMixin, ParametrizedOperator):
     is_binary: bool = True
     has_symmetry_importance: bool = True
@@ -311,35 +145,43 @@ class DateListDiff(PandasOperator, DateDiffMixin, ParametrizedOperator):
             return None
         return cls(aggregation=aggregation)
 
-    def _lists_op(self) -> DateListDiffLists:
-        return DateListDiffLists(
-            diff_unit=self.diff_unit,
-            left_unit=self.left_unit,
-            right_unit=self.right_unit,
-            replace_negative=self.replace_negative,
-        )
-
-    def _agg_op(self) -> DateListDiffAggWithinBounds:
-        return DateListDiffAggWithinBounds(
-            lower_bound=None,
-            upper_bound=None,
-            aggregation=self.aggregation,
-            normalize=False,
-        )
-
-    def _compose_list_diff(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
         if left.isna().all() or right.isna().all():
             return pd.Series([None] * len(left), index=left.index, dtype=np.float64)
 
-        right_mask = DateListDiffLists._non_empty_list_mask(right)
-        diff_lists = self._lists_op().calculate_binary(left, right)
-        result = self._agg_op().calculate_unary(diff_lists)
+        left = self._convert_to_date(left, self.left_unit)
+        right_mask = right.apply(lambda x: len(x) > 0)
+        mask = left.notna() & right.notna() & right_mask
+        right_masked = right[mask].apply(lambda x: pd.arrays.DatetimeArray(self._convert_to_date(x, self.right_unit)))
+
+        if len(right_masked) == 0:
+            diff = []
+        elif len(right_masked) < 2:
+            diff = [left[mask].iloc[0] - right_masked.iloc[0]]
+        else:
+            diff = left[mask] - right_masked.values
+
+        res_masked = pd.Series(diff, index=left[mask].index).apply(lambda x: self._agg(self._diff(x)))
+        res = res_masked.reindex(left.index.union(right.index))
         if self.aggregation in _count_aggregations:
-            result[~right_mask] = 0.0
-        return result.astype(np.float64)
+            res[~right_mask] = 0.0
+        res = res.astype(np.float64)
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
-        return self._compose_list_diff(left, right)
+        return res
+
+    def _diff(self, x: TimedeltaArray):
+        x = self._convert_diff_to_unit(x)
+        return x[x > 0] if self.replace_negative else x
+
+    def _agg(self, x):
+        method = getattr(np, self.aggregation, None)
+        default = np.nan
+        if method is None and self.aggregation in _ext_aggregations:
+            method, default = _ext_aggregations[self.aggregation]
+        elif not callable(method):
+            raise ValueError(f"Unsupported aggregation: {self.aggregation}")
+
+        return method(x) if len(x) > 0 else default
 
 
 class DateListDiffBounded(DateListDiff, ParametrizedOperator):
@@ -388,13 +230,16 @@ class DateListDiffBounded(DateListDiff, ParametrizedOperator):
             res["normalize"] = str(self.normalize)
         return res
 
-    def _agg_op(self) -> DateListDiffAggWithinBounds:
-        return DateListDiffAggWithinBounds(
-            lower_bound=self.lower_bound,
-            upper_bound=self.upper_bound,
-            aggregation=self.aggregation,
-            normalize=self.normalize or False,
-        )
+    def _agg(self, x):
+        orig_len = len(x)
+        x = x[
+            (x >= (self.lower_bound if self.lower_bound is not None else -np.inf))
+            & (x < (self.upper_bound if self.upper_bound is not None else np.inf))
+        ]
+        agg_res = super()._agg(x)
+        if self.normalize and orig_len > 0:
+            return agg_res / orig_len
+        return agg_res
 
 
 class DatePercentileBase(PandasOperator, abc.ABC):
@@ -409,33 +254,26 @@ class DatePercentileBase(PandasOperator, abc.ABC):
         left = pd.to_datetime(left, unit=self.date_unit)
 
         bounds = self._get_bounds(left)
-        values = pd.to_numeric(right, errors="coerce").to_numpy(dtype=np.float64, copy=False)
-        bounds_list = bounds.tolist()
-        result = np.full(len(values), np.nan)
-
-        if not bounds_list:
-            return pd.Series(result, index=right.index).astype(pd.Int64Dtype()).astype("category")
-
-        bounds_lengths = {len(b) for b in bounds_list if isinstance(b, (list, np.ndarray))}
-        if len(bounds_lengths) == 1 and all(isinstance(b, (list, np.ndarray)) for b in bounds_list):
-            bounds_2d = np.asarray(bounds_list, dtype=np.float64)
-            if bounds_2d.ndim == 1:
-                result = bin_index_vectorized(values, bounds_2d)
-            else:
-                result = bin_index_many(values, bounds_2d)
-        else:
-            for i, row_bounds in enumerate(bounds_list):
-                if isinstance(row_bounds, (list, np.ndarray)) and len(row_bounds) > 0:
-                    result[i] = bin_index(values[i], row_bounds)
 
-        return pd.Series(result, index=right.index).astype(pd.Int64Dtype()).astype("category")
+        return (
+            right.index.to_series()
+            .apply(lambda i: self._perc(right[i], bounds[i]))
+            .astype(pd.Int64Dtype())
+            .astype("category")
+        )
 
     @abc.abstractmethod
     def _get_bounds(self, date_col: pd.Series) -> pd.Series:
         pass
 
     def _perc(self, f, bounds):
-        return bin_index(f, bounds)
+        if f is None or np.isnan(f):
+            return np.nan
+        hit = np.where(f >= np.array(bounds))[0]
+        if hit.size > 0:
+            return np.max(hit) + 1
+        else:
+            return np.nan
 
     def get_params(self) -> Dict[str, Optional[str]]:
         res = super().get_params()
@@ -475,15 +313,13 @@ class DatePercentile(DatePercentileBase):
         return value
 
     def _get_bounds(self, date_col: pd.Series) -> pd.Series:
-        zero_bounds = self.zero_bounds if self.zero_bounds is not None else []
-        if not zero_bounds:
-            return pd.Series([[] for _ in range(len(date_col))], index=date_col.index)
-
-        month_diffs = (
-            12 * (date_col.dt.year - (self.zero_year or 0)) + (date_col.dt.month - (self.zero_month or 0))
-        ).to_numpy()
-        bounds_2d = np.asarray(zero_bounds, dtype=np.float64) + month_diffs[:, None] * self.step
-        return pd.Series(list(bounds_2d), index=date_col.index)
+        months = date_col.dt.month
+        years = date_col.dt.year
+
+        month_diffs = 12 * (years - (self.zero_year or 0)) + (months - (self.zero_month or 0))
+        return month_diffs.apply(
+            lambda d: np.array(self.zero_bounds if self.zero_bounds is not None else []) + d * self.step
+        )
 
 
 class DatePercentileMethod2(DatePercentileBase):

{upgini-1.2.156.dev5 → upgini-1.2.157a1.dev1}/src/upgini/autofe/unary.py

@@ -4,7 +4,7 @@ import numpy as np
 import pandas as pd
 
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator, VectorizableMixin
-from upgini.autofe.utils import bin_index, bin_index_vectorized, pydantic_validator
+from upgini.autofe.utils import pydantic_validator
 
 
 class Abs(PandasOperator, VectorizableMixin):
@@ -163,13 +163,16 @@ class Bin(PandasOperator):
     is_categorical: bool = True
 
     def calculate_unary(self, data: pd.Series) -> pd.Series:
-        bounds_arr = np.asarray(self.bin_bounds, dtype=np.float64)
-        values = pd.to_numeric(data, errors="coerce").to_numpy(dtype=np.float64, copy=False)
-        result = bin_index_vectorized(values, bounds_arr)
-        return pd.Series(result, index=data.index).fillna(-1).astype(int).astype("category")
+        return data.apply(self._bin, bounds=self.bin_bounds).fillna(-1).astype(int).astype("category")
 
     def _bin(self, f, bounds):
-        return bin_index(f, bounds)
+        if f is None or np.isnan(f):
+            return np.nan
+        hit = np.where(f >= np.array(bounds))[0]
+        if hit.size > 0:
+            return np.max(hit) + 1
+        else:
+            return np.nan
 
     def get_params(self) -> Dict[str, Optional[str]]:
         res = super().get_params()

{upgini-1.2.156.dev5 → upgini-1.2.157a1.dev1}/src/upgini/autofe/utils.py

@@ -3,9 +3,8 @@ Utility functions for autofe module.
 """
 
 import functools
-from typing import Callable, Union
+from typing import Callable
 
-import numpy as np
 from pydantic import BaseModel
 
 
@@ -112,44 +111,3 @@ def pydantic_copy_method(obj):
         return obj.model_copy
     else:
         return obj.copy
-
-
-def bin_index(value: Union[float, int, None], bounds) -> float:
-    if value is None or (isinstance(value, float) and np.isnan(value)):
-        return np.nan
-    bounds_arr = np.asarray(bounds, dtype=np.float64)
-    if bounds_arr.size == 0 or value < bounds_arr[0]:
-        return np.nan
-    return np.searchsorted(bounds_arr, value, side="right")
-
-
-def bin_index_vectorized(values: np.ndarray, bounds: np.ndarray) -> np.ndarray:
-    n = len(values)
-    result = np.full(n, np.nan)
-    bounds_arr = np.asarray(bounds, dtype=np.float64)
-    if bounds_arr.size == 0:
-        return result
-    valid = ~np.isnan(values)
-    if not valid.any():
-        return result
-    valid_values = values[valid]
-    idx = np.searchsorted(bounds_arr, valid_values, side="right").astype(np.float64)
-    below = valid_values < bounds_arr[0]
-    if below.any():
-        idx[below] = np.nan
-    result[valid] = idx
-    return result
-
-
-def bin_index_many(values: np.ndarray, bounds_2d: np.ndarray) -> np.ndarray:
-    n = len(values)
-    result = np.full(n, np.nan)
-    for i in range(n):
-        v = values[i]
-        if np.isnan(v):
-            continue
-        bounds_row = bounds_2d[i]
-        if bounds_row.size == 0 or v < bounds_row[0]:
-            continue
-        result[i] = np.searchsorted(bounds_row, v, side="right")
-    return result

{upgini-1.2.156.dev5 → upgini-1.2.157a1.dev1}/src/upgini/search_task.py

@@ -1,3 +1,4 @@
+import io
 import logging
 import tempfile
 import time
@@ -438,10 +439,40 @@ def _download_features_file(
     return _read_parquet(file_content)
 
 
+def _is_pyarrow_binary_incompatibility(exc: BaseException) -> bool:
+    message = str(exc)
+    return "binary incompatibility" in message or "IpcReadOptions size changed" in message
+
+
+def _read_parquet_with_pyarrow(file_content: bytes) -> pd.DataFrame:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    # pd.read_parquet(engine="pyarrow") routes through pq.read_table -> ParquetDataset,
+    # which imports pyarrow.dataset and can fail in Colab with ABI mismatch errors.
+    return pq.ParquetFile(pa.BufferReader(file_content)).read().to_pandas()
+
+
+def _read_parquet_with_fastparquet(file_content: bytes) -> pd.DataFrame:
+    import fastparquet
+
+    return fastparquet.ParquetFile(io.BytesIO(file_content)).to_pandas()
+
+
 def _read_parquet(file_content: bytes, file_name: str = "features.parquet"):
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        tmp_file_name = f"{tmp_dir}/{file_name}"
-        with open(tmp_file_name, "wb") as gzip_file:
-            gzip_file.write(file_content)
-        # Note: MLB writes files using pyarrow, so reading with fastparquet may cause errors.
-        return pd.read_parquet(tmp_file_name, engine="pyarrow")
+    try:
+        return _read_parquet_with_pyarrow(file_content)
+    except Exception as pyarrow_error:
+        if not _is_pyarrow_binary_incompatibility(pyarrow_error):
+            raise
+
+    try:
+        # MLB writes files using pyarrow; fastparquet is a fallback for broken pyarrow installs.
+        return _read_parquet_with_fastparquet(file_content)
+    except Exception as fastparquet_error:
+        raise RuntimeError(
+            "Failed to read parquet features file due to PyArrow binary incompatibility. "
+            "If you are using Google Colab, reinstall PyArrow and restart the runtime:\n"
+            "  %pip install --force-reinstall 'pyarrow==23.0.1'\n"
+            "Then use Runtime → Restart session before calling enricher.fit()."
+        ) from fastparquet_error

upgini-1.2.156.dev5/src/upgini/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.2.156.dev5"