Function join_batch_with_features

sdk/python/feast/infra/ray_shared_utils.py:341–462

(batch: pd.DataFrame)

    import ray

    def join_batch_with_features(batch: pd.DataFrame) -> pd.DataFrame:
        features = ray.get(feature_ref)
        if original_join_keys:
            feature_join_keys = original_join_keys
            entity_join_keys = join_keys
        else:
            feature_join_keys = join_keys
            entity_join_keys = join_keys
        feature_cols = [timestamp_field] + feature_join_keys + requested_feats
        available_feature_cols = [
            col for col in feature_cols if col in features.columns
        ]
        features_filtered = features[available_feature_cols].copy()

        batch = normalize_timestamp_columns(batch, timestamp_field, inplace=True)
        features_filtered = normalize_timestamp_columns(
            features_filtered, timestamp_field, inplace=True
        )
        if not entity_join_keys:
            batch_sorted = batch.sort_values(timestamp_field).reset_index(drop=True)
            features_sorted = features_filtered.sort_values(
                timestamp_field
            ).reset_index(drop=True)
            result = pd.merge_asof(
                batch_sorted,
                features_sorted,
                on=timestamp_field,
                direction="backward",
            )
        else:
            for key in entity_join_keys:
                if key not in batch.columns:
                    batch[key] = np.nan
            for key in feature_join_keys:
                if key not in features_filtered.columns:
                    features_filtered[key] = np.nan
            batch_clean = batch.dropna(
                subset=entity_join_keys + [timestamp_field]
            ).copy()
            features_clean = features_filtered.dropna(
                subset=feature_join_keys + [timestamp_field]
            ).copy()
            if batch_clean.empty or features_clean.empty:
                return batch.head(0)
            if timestamp_field in batch_clean.columns:
                batch_sorted = batch_clean.sort_values(
                    timestamp_field, ascending=True
                ).reset_index(drop=True)
            else:
                batch_sorted = batch_clean.reset_index(drop=True)
            right_sort_columns = [
                k for k in feature_join_keys if k in features_clean.columns
            ]
            if timestamp_field in features_clean.columns:
                right_sort_columns.append(timestamp_field)
            if right_sort_columns:
                features_clean = features_clean.drop_duplicates(
                    subset=right_sort_columns, keep="last"