historical feature update

wyvern/feature_store/feature_server.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import importlib
 import logging
+import secrets
 import time
 import traceback
 from collections import defaultdict
@@ -28,11 +29,14 @@
 from feast.value_type import ValueType
 from google.protobuf.json_format import MessageToDict
 
+from wyvern.clients.snowflake import generate_snowflake_ctx
 from wyvern.components.features.realtime_features_component import (
     RealtimeFeatureComponent,
 )
 from wyvern.config import settings
 from wyvern.feature_store.historical_feature_util import (
+    build_and_merge_feast_tables,
+    build_and_merge_realtime_pivot_tables,
     build_historical_real_time_feature_requests,
     build_historical_registry_feature_requests,
     process_historical_real_time_features_requests,
@@ -41,7 +45,9 @@
 )
 from wyvern.feature_store.schemas import (
     GetHistoricalFeaturesRequest,
+    GetHistoricalFeaturesRequestV2,
     GetHistoricalFeaturesResponse,
+    GetHistoricalFeaturesResponseV2,
     GetOnlineFeaturesRequest,
     MaterializeRequest,
 )
@@ -456,6 +462,62 @@ async def get_historical_features(
             results=final_df.to_dict(orient="records"),
         )
 
+    @app.post(f"{settings.WYVERN_HISTORICAL_FEATURES_PATH}_v2")
+    async def get_historical_features_v2(
+        data: GetHistoricalFeaturesRequestV2,
+    ) -> GetHistoricalFeaturesResponseV2:
+        # Generate a 10-digit hex for the request
+        random_id = secrets.token_hex(5)
+
+        # convert the data input to pandas dataframe
+        realtime_features, feast_features = separate_real_time_features(data.features)
+        valid_realtime_features: List[str] = []
+        composite_entities: Dict[str, List[str]] = {}
+        for realtime_feature in realtime_features:
+            entity_type_column = RealtimeFeatureComponent.get_entity_type_column(
+                realtime_feature,
+            )
+            entity_names = RealtimeFeatureComponent.get_entity_names(realtime_feature)
+            if not entity_type_column or not entity_names:
+                logger.warning(f"feature={realtime_feature} is not found")
+                continue
+
+            if len(entity_names) == 2:
+                composite_entities[entity_type_column] = entity_names
+            valid_realtime_features.append(realtime_feature)
+
+        composite_columns = ",".join(
+            [
+                " || ':' || ".join(entities) + f" AS {entity_type_column}"
+                for entity_type_column, entities in composite_entities.items()
+            ],
+        )
+        composite_historical_feature_table = f"HISTORICAL_FEATURES_{random_id}"
+
+        select_sql = f"""
+        CREATE TEMPORARY TABLE {composite_historical_feature_table} AS
+        SELECT *, {composite_columns}, TIMESTAMP as event_timestamp
+        FROM {data.table}
+        """
+        snowflake_ctx = generate_snowflake_ctx()
+        snowflake_ctx.cursor().execute(select_sql)
+
+        result_table = build_and_merge_realtime_pivot_tables(
+            valid_realtime_features,
+            data.table,
+            composite_historical_feature_table,
+            snowflake_ctx,
+        )
+        merged_table = build_and_merge_feast_tables(
+            store,
+            feast_features,
+            result_table,
+            snowflake_ctx,
+        )
+        return GetHistoricalFeaturesResponseV2(
+            result_table=merged_table,
+        )
+
     return app
 
 

wyvern/feature_store/historical_feature_util.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import logging
+import re
 from collections import defaultdict
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Tuple
@@ -117,6 +118,100 @@ def build_historical_real_time_feature_requests(
     return result_dict
 
 
+def build_and_merge_realtime_pivot_tables(
+    full_feature_names: List[str],
+    input_table: str,
+    composite_table: str,
+    context: SnowflakeConnection,
+) -> str:
+    """
+    Build historical real-time feature requests grouped by entity types so that we can process them in parallel.
+
+    Args:
+        full_feature_names: a list of full feature names.
+        request_ids: a list of request ids.
+        entities: a dictionary of entity names and their values.
+
+    Returns:
+        A dictionary of entity types and their corresponding requests.
+    """
+    features_grouped_by_entity = group_realtime_features_by_entity_type(
+        full_feature_names=full_feature_names,
+    )
+    counter = 0
+
+    # prev_table is the previous temporary composite table
+    prev_table = composite_table
+    # next_table is the next temporary composite table joined with the next entity type
+    next_table = f"{composite_table}_0"
+
+    # iterate through all the entity types.
+    # For each entity type, build a new temporary composite table with all the features for this entity type
+    for (
+        entity_identifier_type,
+        curr_feature_names,
+    ) in features_grouped_by_entity.items():
+        entity_list = entity_identifier_type.split(SQL_COLUMN_SEPARATOR)
+
+        if len(entity_list) > 2:
+            logger.warning("Invalid entity_identifier_type={entity_identifier_type}")
+            continue
+        curr_feature_names_underscore = [
+            fn.replace(":", "__", 1) for fn in curr_feature_names
+        ]
+        entity_identifier_type_val = ":".join(entity_list)
+        feature_names_sql_str = ",".join(
+            [f"'{fn}'" for fn in curr_feature_names_underscore],
+        )
+        feature_names_with_pivot_table_str = ",".join(
+            [
+                f"PIVOT_TABLE.{feature_name}"
+                for feature_name in curr_feature_names_underscore
+            ],
+        )
+        feature_names_pivot_raw = ",".join(
+            [f"\"'{fn}'\" as {fn}" for fn in curr_feature_names_underscore],
+        )
+
+        # TODO: send this sql to snowflake
+        pivot_sql = f"""
+        CREATE TEMPORARY TABLE {next_table} AS (
+            WITH PIVOT_DATA AS (
+                SELECT F.REQUEST_ID AS REQUEST,
+                       F.API_SOURCE,
+                       F.EVENT_TYPE,
+                       F.FEATURE_IDENTIFIER,
+                       F.FEATURE_IDENTIFIER_TYPE,
+                       REPLACE(F.FEATURE_NAME, ':', '__') AS FEATURE_NAME,
+                       F.FEATURE_VALUE
+                FROM FEATURE_LOGS F
+                INNER JOIN (SELECT DISTINCT REQUEST FROM {input_table}) T
+                ON F.REQUEST_ID = T.REQUEST
+                WHERE F.FEATURE_IDENTIFIER_TYPE = '{entity_identifier_type_val}'
+            ), PIVOT_TABLE_RAW AS (
+                SELECT *
+                FROM PIVOT_DATA
+                PIVOT(MAX(FEATURE_VALUE) FOR FEATURE_NAME IN ({feature_names_sql_str}))
+            ), PIVOT_TABLE AS (
+                SELECT REQUEST, FEATURE_IDENTIFIER, FEATURE_IDENTIFIER_TYPE, {feature_names_pivot_raw}
+                FROM PIVOT_TABLE_RAW
+            )
+            SELECT
+                {prev_table}.*,{feature_names_with_pivot_table_str}
+            FROM
+                {prev_table}
+            LEFT JOIN PIVOT_TABLE ON
+                {prev_table}.REQUEST = PIVOT_TABLE.REQUEST AND
+                {prev_table}.{entity_identifier_type} = PIVOT_TABLE.FEATURE_IDENTIFIER
+        )
+        """
+        context.cursor().execute(pivot_sql)
+        counter += 1
+        prev_table = next_table
+        next_table = f"{composite_table}_{counter}"
+    return prev_table
+
+
 def process_historical_real_time_features_requests(
     requests: Dict[str, RequestEntityIdentifierObjects],
 ) -> Dict[str, pd.DataFrame]:
@@ -316,6 +411,109 @@ def build_historical_registry_feature_requests(
     return requests
 
 
+def build_and_merge_feast_tables(
+    store: FeatureStore,
+    feature_names: List[str],
+    composite_table: str,
+    context: SnowflakeConnection,
+) -> str:
+    features_grouped_by_entities = group_registry_features_by_entities(
+        feature_names,
+        store=store,
+    )
+    counter = 0
+    prev_table = composite_table
+    next_table = f"{composite_table}_0"
+    for entity_name, feature_names in features_grouped_by_entities.items():
+        if not feature_names:
+            continue
+
+        if FULL_FEATURE_NAME_SEPARATOR in entity_name:
+            entities = entity_name.split(FULL_FEATURE_NAME_SEPARATOR)
+        else:
+            entities = [entity_name]
+
+        if len(entities) > 2:
+            raise ValueError(
+                f"Entity name should be singular or composite: {entity_name}",
+            )
+
+        feature_columns = [fn.replace(":", "__") for fn in feature_names]
+
+        # TODO: validate that all entities are in the entity_df_table
+        # for entity in entities:
+        # if entity not in entity_values:
+        #     raise ValueError(
+        #         f"{feature_names} depends on {entity}. Could not find entity values: {entity}",
+        #     )
+        identifier_column = SQL_COLUMN_SEPARATOR.join(entities)
+        identifier_table_sql_dupe = f"""
+        SELECT
+            {identifier_column} AS IDENTIFIER,
+            event_timestamp,
+            ROW_NUMBER() OVER (PARTITION BY IDENTIFIER, event_timestamp ORDER BY (SELECT NULL)) as rn
+        FROM {composite_table}
+        WHERE {identifier_column} is NOT NULL
+        """
+
+        # dedupe (IDENTIFIER, event_timestamp)
+        identifier_table_sql = f"""
+        WITH identifier_table_sql_dupe AS ({identifier_table_sql_dupe})
+        SELECT IDENTIFIER, event_timestamp
+        FROM identifier_table_sql_dupe
+        WHERE rn = 1
+        """
+        result = store.get_historical_features(
+            entity_df=identifier_table_sql,
+            features=feature_names or [],
+            full_feature_names=True,
+        )
+        result_sql = result.to_sql()
+        # Strip the leading "WITH " (WITH plus an empty space)
+        result_sql = result_sql.replace("WITH ", "")
+        # Replace the table name with 'identifier_tbl', assuming the table name is always
+        # in the format "feast_entity_df_" followed by a hex string (UUID without dashes)
+        result_sql = re.sub(
+            r'"feast_entity_df_[0-9a-f]{32}"',
+            "identifier_tbl",
+            result_sql,
+            flags=re.IGNORECASE,
+        )
+        new_feast_table_sql = f"""
+        CREATE TEMPORARY TABLE {next_table}_feast AS (
+        WITH identifier_tbl_dupe AS ({identifier_table_sql_dupe}),
+        identifier_tbl AS (
+            SELECT IDENTIFIER, event_timestamp
+            FROM identifier_tbl_dupe
+            WHERE rn = 1
+        ),
+            {result_sql}
+        )
+        """
+        context.cursor().execute(new_feast_table_sql)
+
+        # left join to the previous composite table
+        picked_feature_columns_str = ", ".join(
+            [f'{next_table}_feast."{c}"' for c in feature_columns],
+        )
+        new_composite_table_sql = f"""
+        CREATE TABLE {next_table} AS (
+            SELECT {prev_table}.*, {picked_feature_columns_str}
+            FROM {prev_table}
+            LEFT JOIN {next_table}_feast
+            ON {prev_table}.{identifier_column} = {next_table}_feast.IDENTIFIER and
+                {prev_table}.event_timestamp = {next_table}_feast.event_timestamp
+        )
+        """
+        context.cursor().execute(new_composite_table_sql)
+
+        counter += 1
+        prev_table = next_table
+        next_table = f"{composite_table}_{counter}"
+
+    return prev_table
+
+
 def process_historical_registry_features_requests(
     store: FeatureStore,
     requests: List[GetFeastHistoricalFeaturesRequest],
@@ -354,12 +552,18 @@ def process_historical_registry_features_request(
     entity_df = pd.DataFrame(request.entities)
     # no timezone is allowed in the timestamp
     entity_df["event_timestamp"] = entity_df["event_timestamp"].dt.tz_localize(None)
+    # TODO: use sql to get the result.
+    # example:
+    # https://docs.feast.dev/getting-started/concepts/feature-retrieval
+    # #example-entity-sql-query-for-generating-training-data
     result = store.get_historical_features(
         entity_df=entity_df,
         features=request.features or [],
         full_feature_names=request.full_feature_names,
     )
+    # TODO: to_sql(); replace IDENTIFIER by the original identifier_type
     result_df = result.to_df()
+    # TODO: group IDENTIFIER and event_timestamp
     result_df.drop_duplicates(subset=["IDENTIFIER", "event_timestamp"], inplace=True)
     return entity_df.merge(
         result_df,

wyvern/feature_store/schemas.py

@@ -38,6 +38,21 @@ class GetHistoricalFeaturesRequest(BaseModel):
     features: List[str] = []
 
 
+class GetHistoricalFeaturesRequestV2(BaseModel):
+    """
+    Request object for getting historical features.
+
+    Attributes:
+        entities: A dictionary of entity name to entity value.
+        timestamps: A list of timestamps. Used to retrieve historical features at specific timestamps. If not provided,
+            the latest feature values will be returned.
+        features: A list of feature names.
+    """
+
+    table: str
+    features: List[str] = []
+
+
 class GetFeastHistoricalFeaturesRequest(BaseModel):
     """
     Request object for getting historical features from Feast.
@@ -66,6 +81,17 @@ class GetHistoricalFeaturesResponse(BaseModel):
     results: List[Dict[str, Any]] = []
 
 
+class GetHistoricalFeaturesResponseV2(BaseModel):
+    """
+    Response object for getting historical features.
+
+    Attributes:
+        result_table: the name of the transient table that holds the result.
+    """
+
+    result_table: str
+
+
 class MaterializeRequest(BaseModel):
     """
     Request object for materializing feature views.