Implement entity aliasing for feature retrival (feast-dev#1868)

* Initial commit for implementing shadow entities 8b

Signed-off-by: Cody Lin <[email protected]>

* fix small bug failing unit tests

Signed-off-by: Cody Lin <[email protected]>

* revert changes using FVProjection

Signed-off-by: Cody Lin <[email protected]>

* initial design for testing

Signed-off-by: Cody Lin <[email protected]>

* add join_key_map to FV proto definition

Signed-off-by: Cody Lin <[email protected]>

* handle join_key_map for proto conversions

Signed-off-by: Cody Lin <[email protected]>

* try to grab underlying changes

Signed-off-by: Cody Lin <[email protected]>

* Tests fixed after rebasing with dependencies

Signed-off-by: Cody Lin <[email protected]>

* Remove small unnecessary changes

Signed-off-by: Cody Lin <[email protected]>

* Remove small unnecessary changes

Signed-off-by: Cody Lin <[email protected]>

* Fix lint issues

Signed-off-by: Cody Lin <[email protected]>

* Fix redshift query error

Signed-off-by: Cody Lin <[email protected]>

* Fix cache error

Signed-off-by: Cody Lin <[email protected]>

* Fix benchmark test error

Signed-off-by: Cody Lin <[email protected]>

* remove unnecessary changes

Signed-off-by: Cody Lin <[email protected]>

* respond to David's comments

Signed-off-by: Cody Lin <[email protected]>

* Improve docstring, revert to entity_name_to_join_key_map

Signed-off-by: Cody Lin <[email protected]>

* Fix docstring test

Signed-off-by: Cody Lin <[email protected]>

* Remove code from previous design

Signed-off-by: Cody Lin <[email protected]>

* Improve comments and variable naming

Signed-off-by: Cody Lin <[email protected]>

* Use get with default

Signed-off-by: Cody Lin <[email protected]>

* Fix use get with default

Signed-off-by: Cody Lin <[email protected]>

Co-authored-by: David Y Liu <[email protected]>

Author: codyjlin

protos/feast/core/FeatureViewProjection.proto

@@ -19,4 +19,7 @@ message FeatureViewProjection {
 
   // The features of the feature view that are a part of the feature reference.
   repeated FeatureSpecV2 feature_columns = 2;
+
+  // Map for entity join_key overrides of feature data entity join_key to entity data join_key
+  map<string,string> join_key_map = 4;
 }

sdk/python/feast/driver_test_data.py

@@ -1,4 +1,5 @@
 # This module generates dummy data to be used for tests and examples.
+import itertools
 from enum import Enum
 
 import numpy as np
@@ -29,22 +30,29 @@ def _convert_event_timestamp(event_timestamp: pd.Timestamp, t: EventTimestampTyp
 
 
 def create_orders_df(
-    customers, drivers, start_date, end_date, order_count,
+    customers, drivers, start_date, end_date, order_count, locations=None,
 ) -> pd.DataFrame:
     """
-    Example df generated by this function:
+    Example df generated by this function (if locations):
 
-    | order_id | driver_id | customer_id | order_is_success |    event_timestamp  |
-    +----------+-----------+-------------+------------------+---------------------+
-    |      100 |      5004 |        1007 |                0 | 2021-03-10 19:31:15 |
-    |      101 |      5003 |        1006 |                0 | 2021-03-11 22:02:50 |
-    |      102 |      5010 |        1005 |                0 | 2021-03-13 00:34:24 |
-    |      103 |      5010 |        1001 |                1 | 2021-03-14 03:05:59 |
+    | order_id | driver_id | customer_id | origin_id | destination_id | order_is_success | event_timestamp     |
+    +----------+-----------+-------------+-----------+----------------+------------------+---------------------+
+    |      100 |      5004 |        1007 |         1 |             18 |                0 | 2021-03-10 19:31:15 |
+    |      101 |      5003 |        1006 |        24 |             42 |                0 | 2021-03-11 22:02:50 |
+    |      102 |      5010 |        1005 |        19 |             12 |                0 | 2021-03-13 00:34:24 |
+    |      103 |      5010 |        1001 |        35 |              8 |                1 | 2021-03-14 03:05:59 |
     """
     df = pd.DataFrame()
     df["order_id"] = [order_id for order_id in range(100, 100 + order_count)]
     df["driver_id"] = np.random.choice(drivers, order_count)
     df["customer_id"] = np.random.choice(customers, order_count)
+    if locations:
+        location_pairs = np.array(list(itertools.permutations(locations, 2)))
+        locations_sample = location_pairs[
+            np.random.choice(len(location_pairs), order_count)
+        ].T
+        df["origin_id"] = locations_sample[0]
+        df["destination_id"] = locations_sample[1]
     df["order_is_success"] = np.random.randint(0, 2, size=order_count).astype(np.int32)
     df[DEFAULT_ENTITY_DF_EVENT_TIMESTAMP_COL] = [
         _convert_event_timestamp(
@@ -180,6 +188,46 @@ def create_customer_daily_profile_df(customers, start_date, end_date) -> pd.Data
     return df_all_customers
 
 
+def create_location_stats_df(locations, start_date, end_date) -> pd.DataFrame:
+    """
+    Example df generated by this function:
+
+    | event_timestamp  | location_id | temperature | created          |
+    +------------------+-------------+-------------+------------------+
+    | 2021-03-17 19:31 |           1 |          74 | 2021-03-24 19:38 |
+    | 2021-03-17 20:31 |          24 |          63 | 2021-03-24 19:38 |
+    | 2021-03-17 21:31 |          19 |          65 | 2021-03-24 19:38 |
+    | 2021-03-17 22:31 |          35 |          86 | 2021-03-24 19:38 |
+    """
+    df_hourly = pd.DataFrame(
+        {
+            "event_timestamp": [
+                pd.Timestamp(dt, unit="ms", tz="UTC").round("ms")
+                for dt in pd.date_range(
+                    start=start_date, end=end_date, freq="1H", closed="left"
+                )
+            ]
+        }
+    )
+    df_all_locations = pd.DataFrame()
+
+    for location in locations:
+        df_hourly_copy = df_hourly.copy()
+        df_hourly_copy["location_id"] = location
+        df_all_locations = pd.concat([df_hourly_copy, df_all_locations])
+
+    df_all_locations.reset_index(drop=True, inplace=True)
+    rows = df_all_locations["event_timestamp"].count()
+
+    df_all_locations["temperature"] = np.random.randint(50, 100, size=rows).astype(
+        np.int32
+    )
+
+    # TODO: Remove created timestamp in order to test whether its really optional
+    df_all_locations["created"] = pd.to_datetime(pd.Timestamp.now(tz=None).round("ms"))
+    return df_all_locations
+
+
 def create_global_daily_stats_df(start_date, end_date) -> pd.DataFrame:
     """
     Example df generated by this function:

sdk/python/feast/feature_store.py

@@ -839,6 +839,20 @@ def get_online_features(
         entity_name_to_join_key_map = {}
         for entity in entities:
             entity_name_to_join_key_map[entity.name] = entity.join_key
+        for feature_view in all_feature_views:
+            for entity_name in feature_view.entities:
+                entity = self._registry.get_entity(
+                    entity_name, self.project, allow_cache=True
+                )
+                # User directly uses join_key as the entity reference in the entity_rows for the
+                # entity mapping case.
+                entity_name = feature_view.projection.join_key_map.get(
+                    entity.join_key, entity.name
+                )
+                join_key = feature_view.projection.join_key_map.get(
+                    entity.join_key, entity.join_key
+                )
+                entity_name_to_join_key_map[entity_name] = join_key
 
         needed_request_data_features = self._get_needed_request_data_features(
             grouped_odfv_refs
@@ -895,8 +909,12 @@ def get_online_features(
                 ] = GetOnlineFeaturesResponse.FieldStatus.PRESENT
 
         for table, requested_features in grouped_refs:
+            table_join_keys = [
+                entity_name_to_join_key_map[entity_name]
+                for entity_name in table.entities
+            ]
             self._populate_result_rows_from_feature_view(
-                entity_name_to_join_key_map,
+                table_join_keys,
                 full_feature_names,
                 provider,
                 requested_features,
@@ -918,7 +936,7 @@ def get_online_features(
 
     def _populate_result_rows_from_feature_view(
         self,
-        entity_name_to_join_key_map: Dict[str, str],
+        table_join_keys: List[str],
         full_feature_names: bool,
         provider: Provider,
         requested_features: List[str],
@@ -927,7 +945,7 @@ def _populate_result_rows_from_feature_view(
         union_of_entity_keys: List[EntityKeyProto],
     ):
         entity_keys = _get_table_entity_keys(
-            table, union_of_entity_keys, entity_name_to_join_key_map
+            table, union_of_entity_keys, table_join_keys
         )
         read_rows = provider.online_read(
             config=self.config,
@@ -1045,8 +1063,8 @@ def _get_feature_views_to_use(
             )
         }
 
-        fvs_to_use, od_fvs_to_use = [], []
         if isinstance(features, FeatureService):
+            fvs_to_use, od_fvs_to_use = [], []
             for fv_name, projection in [
                 (projection.name, projection)
                 for projection in features.feature_view_projections
@@ -1137,10 +1155,12 @@ def _group_feature_refs(
     """ Get list of feature views and corresponding feature names based on feature references"""
 
     # view name to view proto
-    view_index = {view.name: view for view in all_feature_views}
+    view_index = {view.projection.name_to_use(): view for view in all_feature_views}
 
     # on demand view to on demand view proto
-    on_demand_view_index = {view.name: view for view in all_on_demand_feature_views}
+    on_demand_view_index = {
+        view.projection.name_to_use(): view for view in all_on_demand_feature_views
+    }
 
     # view name to feature names
     views_features = defaultdict(list)
@@ -1168,15 +1188,19 @@ def _group_feature_refs(
 
 
 def _get_table_entity_keys(
-    table: FeatureView, entity_keys: List[EntityKeyProto], join_key_map: Dict[str, str],
+    table: FeatureView, entity_keys: List[EntityKeyProto], table_join_keys: List[str]
 ) -> List[EntityKeyProto]:
-    table_join_keys = [join_key_map[entity_name] for entity_name in table.entities]
+    reverse_join_key_map = {
+        alias: original for original, alias in table.projection.join_key_map.items()
+    }
     required_entities = OrderedDict.fromkeys(sorted(table_join_keys))
     entity_key_protos = []
     for entity_key in entity_keys:
         required_entities_to_values = required_entities.copy()
         for i in range(len(entity_key.join_keys)):
-            entity_name = entity_key.join_keys[i]
+            entity_name = reverse_join_key_map.get(
+                entity_key.join_keys[i], entity_key.join_keys[i]
+            )
             entity_value = entity_key.entity_values[i]
 
             if entity_name in required_entities_to_values:

sdk/python/feast/feature_view.py

@@ -126,7 +126,7 @@ def __init__(
         self.name = name
         self.entities = entities if entities else [DUMMY_ENTITY_NAME]
         self.features = _features
-        self.tags = tags if tags is not None else {}
+        self.tags = tags if tags else {}
 
         if isinstance(ttl, Duration):
             self.ttl = timedelta(seconds=int(ttl.seconds))
@@ -238,6 +238,44 @@ def with_name(self, name: str):
 
         return cp
 
+    def with_join_key_map(self, join_key_map: Dict[str, str]):
+        """
+        Sets the join_key_map by returning a copy of this feature view with that field set.
+        This join_key mapping operation is only used as part of query operations and will
+        not modify the underlying FeatureView.
+
+        Args:
+            join_key_map: A map of join keys in which the left is the join_key that
+                corresponds with the feature data and the right corresponds with the entity data.
+
+        Returns:
+            A copy of this FeatureView with the join_key_map replaced with the 'join_key_map' input.
+
+        Examples:
+            Join a location feature data table to both the origin column and destination
+            column of the entity data.
+
+            temperatures_feature_service = FeatureService(
+                name="temperatures",
+                features=[
+                    location_stats_feature_view
+                        .with_name("origin_stats")
+                        .with_join_key_map(
+                            {"location_id": "origin_id"}
+                        ),
+                    location_stats_feature_view
+                        .with_name("destination_stats")
+                        .with_join_key_map(
+                            {"location_id": "destination_id"}
+                        ),
+                ],
+            )
+        """
+        cp = self.__copy__()
+        cp.projection.join_key_map = join_key_map
+
+        return cp
+
     def with_projection(self, feature_view_projection: FeatureViewProjection):
         """
         Sets the feature view projection by returning a copy of this feature view

sdk/python/feast/feature_view_projection.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from attr import dataclass
 
@@ -13,13 +13,16 @@ class FeatureViewProjection:
     name: str
     name_alias: Optional[str]
     features: List[Feature]
+    join_key_map: Dict[str, str] = {}
 
     def name_to_use(self):
         return self.name_alias or self.name
 
     def to_proto(self):
         feature_reference_proto = FeatureViewProjectionProto(
-            feature_view_name=self.name, feature_view_name_alias=self.name_alias
+            feature_view_name=self.name,
+            feature_view_name_alias=self.name_alias,
+            join_key_map=self.join_key_map,
         )
         for feature in self.features:
             feature_reference_proto.feature_columns.append(feature.to_proto())
@@ -32,6 +35,7 @@ def from_proto(proto: FeatureViewProjectionProto):
             name=proto.feature_view_name,
             name_alias=proto.feature_view_name_alias,
             features=[],
+            join_key_map=dict(proto.join_key_map),
         )
         for feature_column in proto.feature_columns:
             ref.features.append(Feature.from_proto(feature_column))

sdk/python/feast/infra/offline_stores/file.py

@@ -143,6 +143,11 @@ def evaluate_historical_retrieval():
                     table = _run_field_mapping(
                         table, feature_view.batch_source.field_mapping
                     )
+                # Rename entity columns by the join_key_map dictionary if it exists
+                if feature_view.projection.join_key_map:
+                    table = _run_field_mapping(
+                        table, feature_view.projection.join_key_map
+                    )
 
                 # Convert pyarrow table to pandas dataframe. Note, if the underlying data has missing values,
                 # pandas will convert those values to np.nan if the dtypes are numerical (floats, ints, etc.) or boolean
@@ -176,7 +181,9 @@ def evaluate_historical_retrieval():
                     # double underscore as separator for consistency with other databases like BigQuery,
                     # where there are very few characters available for use as separators
                     if full_feature_names:
-                        formatted_feature_name = f"{feature_view.name}__{feature}"
+                        formatted_feature_name = (
+                            f"{feature_view.projection.name_to_use()}__{feature}"
+                        )
                     else:
                         formatted_feature_name = feature
                     # Add the feature name to the list of columns
@@ -191,7 +198,10 @@ def evaluate_historical_retrieval():
                 join_keys = []
                 for entity_name in feature_view.entities:
                     entity = registry.get_entity(entity_name, project)
-                    join_keys.append(entity.join_key)
+                    join_key = feature_view.projection.join_key_map.get(
+                        entity.join_key, entity.join_key
+                    )
+                    join_keys.append(join_key)
                 right_entity_columns = join_keys
                 right_entity_key_columns = [
                     event_timestamp_column

sdk/python/feast/infra/offline_stores/offline_utils.py

@@ -64,7 +64,10 @@ def get_expected_join_keys(
         entities = feature_view.entities
         for entity_name in entities:
             entity = registry.get_entity(entity_name, project)
-            join_keys.add(entity.join_key)
+            join_key = feature_view.projection.join_key_map.get(
+                entity.join_key, entity.join_key
+            )
+            join_keys.add(join_key)
     return join_keys
 
 
@@ -113,11 +116,11 @@ def get_feature_view_query_context(
         }
         for entity_name in feature_view.entities:
             entity = registry.get_entity(entity_name, project)
-            join_keys.append(entity.join_key)
-            join_key_column = reverse_field_mapping.get(
+            join_key = feature_view.projection.join_key_map.get(
                 entity.join_key, entity.join_key
             )
-            entity_selections.append(f"{join_key_column} AS {entity.join_key}")
+            join_keys.append(join_key)
+            entity_selections.append(f"{entity.join_key} AS {join_key}")
 
         if isinstance(feature_view.ttl, timedelta):
             ttl_seconds = int(feature_view.ttl.total_seconds())

sdk/python/feast/infra/offline_stores/redshift.py

@@ -176,7 +176,7 @@ def query_generator() -> Iterator[str]:
             ),
             drop_columns=["entity_timestamp"]
             + [
-                f"{feature_view.name}__entity_row_unique_id"
+                f"{feature_view.projection.name_to_use()}__entity_row_unique_id"
                 for feature_view in feature_views
             ],
         )

sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py

@@ -6,7 +6,11 @@
 from tests.integration.feature_repos.repo_configuration import (
     construct_universal_feature_views,
 )
-from tests.integration.feature_repos.universal.entities import customer, driver
+from tests.integration.feature_repos.universal.entities import (
+    customer,
+    driver,
+    location,
+)
 
 
 @pytest.mark.benchmark
@@ -24,7 +28,7 @@ def test_online_retrieval(environment, universal_data_sources, benchmark):
 
     feast_objects = []
     feast_objects.extend(feature_views.values())
-    feast_objects.extend([driver(), customer(), feature_service])
+    feast_objects.extend([driver(), customer(), location(), feature_service])
     fs.apply(feast_objects)
     fs.materialize(environment.start_date, environment.end_date)