Repository: feast-dev/feast

Issue:
Redis can delete incorrect keys
I haven't tested this explicitly but from the data model + this function it looks to be as if two FeatureViews have the same Entities (but are different tables) and one view is deleted the online keys for the other view will also be deleted... This seems like a bit of a flaw either in the Redis data model or in the deletion mechanism. 

https://github.com/feast-dev/feast/blob/ec4165396f70ab20b42246b093f777dfcc9f5277/sdk/python/feast/infra/online_stores/redis.py#L75-L88


Discussion:
good catch - briefly looking at the code, I'm not even convinced that the key we're using to delete matches the key for insertion (e.g. we might be deleting nothing at all)

Resolving patch:
diff --git a/sdk/python/feast/infra/online_stores/redis.py b/sdk/python/feast/infra/online_stores/redis.py
--- a/sdk/python/feast/infra/online_stores/redis.py
+++ b/sdk/python/feast/infra/online_stores/redis.py
@@ -72,11 +72,11 @@ class RedisOnlineStoreConfig(FeastConfigBaseModel):
 class RedisOnlineStore(OnlineStore):
     _client: Optional[Union[Redis, RedisCluster]] = None
 
-    def delete_table_values(self, config: RepoConfig, table: FeatureView):
+    def delete_entity_values(self, config: RepoConfig, join_keys: List[str]):
         client = self._get_client(config.online_store)
         deleted_count = 0
         pipeline = client.pipeline()
-        prefix = _redis_key_prefix(table.entities)
+        prefix = _redis_key_prefix(join_keys)
 
         for _k in client.scan_iter(
             b"".join([prefix, b"*", config.project.encode("utf8")])
@@ -85,7 +85,7 @@ def delete_table_values(self, config: RepoConfig, table: FeatureView):
             deleted_count += 1
         pipeline.execute()
 
-        logger.debug(f"Deleted {deleted_count} keys for {table.name}")
+        logger.debug(f"Deleted {deleted_count} rows for entity {', '.join(join_keys)}")
 
     @log_exceptions_and_usage(online_store="redis")
     def update(
@@ -98,10 +98,16 @@ def update(
         partial: bool,
     ):
         """
-        We delete the keys in redis for tables/views being removed.
+        Look for join_keys (list of entities) that are not in use anymore
+        (usually this happens when the last feature view that was using specific compound key is deleted)
+        and remove all features attached to this "join_keys".
         """
-        for table in tables_to_delete:
-            self.delete_table_values(config, table)
+        join_keys_to_keep = set(tuple(table.entities) for table in tables_to_keep)
+
+        join_keys_to_delete = set(tuple(table.entities) for table in tables_to_delete)
+
+        for join_keys in join_keys_to_delete - join_keys_to_keep:
+            self.delete_entity_values(config, list(join_keys))
 
     def teardown(
         self,
@@ -112,8 +118,10 @@ def teardown(
         """
         We delete the keys in redis for tables/views being removed.
         """
-        for table in tables:
-            self.delete_table_values(config, table)
+        join_keys_to_delete = set(tuple(table.entities) for table in tables)
+
+        for join_keys in join_keys_to_delete:
+            self.delete_entity_values(config, list(join_keys))
 
     @staticmethod
     def _parse_connection_string(connection_string: str):


Test patch:
diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py
--- a/sdk/python/tests/integration/feature_repos/repo_configuration.py
+++ b/sdk/python/tests/integration/feature_repos/repo_configuration.py
@@ -7,7 +7,7 @@
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import pandas as pd
 import yaml
@@ -283,7 +283,9 @@ def construct_test_environment(
             execution_role_name="arn:aws:iam::402087665549:role/lambda_execution_role",
         )
 
-        registry = f"s3://feast-integration-tests/registries/{project}/registry.db"
+        registry = (
+            f"s3://feast-integration-tests/registries/{project}/registry.db"
+        )  # type: Union[str, RegistryConfig]
     else:
         # Note: even if it's a local feature server, the repo config does not have this configured
         feature_server = None
diff --git a/sdk/python/tests/integration/online_store/test_universal_online.py b/sdk/python/tests/integration/online_store/test_universal_online.py
--- a/sdk/python/tests/integration/online_store/test_universal_online.py
+++ b/sdk/python/tests/integration/online_store/test_universal_online.py
@@ -28,6 +28,7 @@
 )
 from tests.integration.feature_repos.universal.feature_views import (
     create_driver_hourly_stats_feature_view,
+    driver_feature_view,
 )
 from tests.utils.data_source_utils import prep_file_source
 
@@ -503,6 +504,79 @@ def test_online_retrieval(environment, universal_data_sources, full_feature_name
     )
 
 
+@pytest.mark.integration
+@pytest.mark.universal
+def test_online_store_cleanup(environment, universal_data_sources):
+    """
+    Some online store implementations (like Redis) keep features from different features views
+    but with common entities together.
+    This might end up with deletion of all features attached to the entity,
+    when only one feature view was deletion target (see https://github.com/feast-dev/feast/issues/2150).
+
+    Plan:
+        1. Register two feature views with common entity "driver"
+        2. Materialize data
+        3. Check if features are available (via online retrieval)
+        4. Delete one feature view
+        5. Check that features for other are still available
+        6. Delete another feature view (and create again)
+        7. Verify that features for both feature view were deleted
+    """
+    fs = environment.feature_store
+    entities, datasets, data_sources = universal_data_sources
+    driver_stats_fv = construct_universal_feature_views(data_sources)["driver"]
+
+    df = pd.DataFrame(
+        {
+            "ts_1": [environment.end_date] * len(entities["driver"]),
+            "created_ts": [environment.end_date] * len(entities["driver"]),
+            "driver_id": entities["driver"],
+            "value": np.random.random(size=len(entities["driver"])),
+        }
+    )
+
+    ds = environment.data_source_creator.create_data_source(
+        df, destination_name="simple_driver_dataset"
+    )
+
+    simple_driver_fv = driver_feature_view(
+        data_source=ds, name="test_universal_online_simple_driver"
+    )
+
+    fs.apply([driver(), simple_driver_fv, driver_stats_fv])
+
+    fs.materialize(
+        environment.start_date - timedelta(days=1),
+        environment.end_date + timedelta(days=1),
+    )
+    expected_values = df.sort_values(by="driver_id")
+
+    features = [f"{simple_driver_fv.name}:value"]
+    entity_rows = [{"driver": driver_id} for driver_id in sorted(entities["driver"])]
+
+    online_features = fs.get_online_features(
+        features=features, entity_rows=entity_rows
+    ).to_dict()
+    assert np.allclose(expected_values["value"], online_features["value"])
+
+    fs.apply(
+        objects=[simple_driver_fv], objects_to_delete=[driver_stats_fv], partial=False
+    )
+
+    online_features = fs.get_online_features(
+        features=features, entity_rows=entity_rows
+    ).to_dict()
+    assert np.allclose(expected_values["value"], online_features["value"])
+
+    fs.apply(objects=[], objects_to_delete=[simple_driver_fv], partial=False)
+    fs.apply([simple_driver_fv])
+
+    online_features = fs.get_online_features(
+        features=features, entity_rows=entity_rows
+    ).to_dict()
+    assert all(v is None for v in online_features["value"])
+
+
 def response_feature_name(feature: str, full_feature_names: bool) -> str:
     if (
         feature in {"current_balance", "avg_passenger_count", "lifetime_trip_count"}