Delete keys from Redis when tearing down online store (feast-dev#1965)

achals · web-flow · commit 7538f7e64376 · 2021-10-23T23:17:45.000-07:00
* Delete keys from redis when tearing down online store

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* add project

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Update comments and update method as well

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Update comments and update method as well

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* CR updates

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Update comments and update method as well

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* docstring

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/feature_table.py b/sdk/python/feast/feature_table.py
@@ -110,7 +110,7 @@ def name(self, name: str):
         self._name = name
 
     @property
-    def entities(self):
+    def entities(self) -> List[str]:
         """
         Returns the entities of this feature table
         """
diff --git a/sdk/python/feast/infra/key_encoding_utils.py b/sdk/python/feast/infra/key_encoding_utils.py
@@ -19,6 +19,22 @@ def _serialize_val(value_type, v: ValueProto) -> Tuple[bytes, int]:
         raise ValueError(f"Value type not supported for Firestore: {v}")
 
 
+def serialize_entity_key_prefix(entity_keys: List[str]) -> bytes:
+    """
+    Serialize keys to a bytestring so it can be used to prefix-scan through items stored in the online store
+    using serialize_entity_key.
+
+    This encoding is a partial implementation of serialize_entity_key, only operating on the keys of entities,
+    and not the values.
+    """
+    sorted_keys = sorted(entity_keys)
+    output: List[bytes] = []
+    for k in sorted_keys:
+        output.append(struct.pack("<I", ValueType.STRING))
+        output.append(k.encode("utf8"))
+    return b"".join(output)
+
+
 def serialize_entity_key(entity_key: EntityKeyProto) -> bytes:
     """
     Serialize entity key to a bytestring so it can be used as a lookup key in a hash table.
diff --git a/sdk/python/feast/infra/online_stores/helpers.py b/sdk/python/feast/infra/online_stores/helpers.py
@@ -5,7 +5,10 @@
 import mmh3
 
 from feast import errors
-from feast.infra.key_encoding_utils import serialize_entity_key
+from feast.infra.key_encoding_utils import (
+    serialize_entity_key,
+    serialize_entity_key_prefix,
+)
 from feast.infra.online_stores.online_store import OnlineStore
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 
@@ -41,6 +44,10 @@ def _redis_key(project: str, entity_key: EntityKeyProto) -> bytes:
     return b"".join(key)
 
 
+def _redis_key_prefix(entity_keys: List[str]) -> bytes:
+    return serialize_entity_key_prefix(entity_keys)
+
+
 def _mmh3(key: str):
     """
     Calculate murmur3_32 hash which is equal to scala version which is using little endian:
diff --git a/sdk/python/feast/infra/online_stores/redis.py b/sdk/python/feast/infra/online_stores/redis.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import logging
 from datetime import datetime
 from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -21,7 +22,7 @@
 from pydantic.typing import Literal
 
 from feast import Entity, FeatureTable, FeatureView, RepoConfig, utils
-from feast.infra.online_stores.helpers import _mmh3, _redis_key
+from feast.infra.online_stores.helpers import _mmh3, _redis_key, _redis_key_prefix
 from feast.infra.online_stores.online_store import OnlineStore
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
@@ -36,6 +37,7 @@
     raise FeastExtrasDependencyImportError("redis", str(e))
 
 EX_SECONDS = 253402300799
+logger = logging.getLogger(__name__)
 
 
 class RedisType(str, Enum):
@@ -60,6 +62,23 @@ class RedisOnlineStoreConfig(FeastConfigBaseModel):
 class RedisOnlineStore(OnlineStore):
     _client: Optional[Union[Redis, RedisCluster]] = None
 
+    def delete_table_values(
+        self, config: RepoConfig, table: Union[FeatureTable, FeatureView]
+    ):
+        client = self._get_client(config.online_store)
+        deleted_count = 0
+        pipeline = client.pipeline()
+        prefix = _redis_key_prefix(table.entities)
+
+        for _k in client.scan_iter(
+            b"".join([prefix, b"*", config.project.encode("utf8")])
+        ):
+            pipeline.delete(_k)
+            deleted_count += 1
+        pipeline.execute()
+
+        logger.debug(f"Deleted {deleted_count} keys for {table.name}")
+
     def update(
         self,
         config: RepoConfig,
@@ -70,9 +89,10 @@ def update(
         partial: bool,
     ):
         """
-        There's currently no setup done for Redis.
+        We delete the keys in redis for tables/views being removed.
         """
-        pass
+        for table in tables_to_delete:
+            self.delete_table_values(config, table)
 
     def teardown(
         self,
@@ -81,9 +101,10 @@ def teardown(
         entities: Sequence[Entity],
     ):
         """
-        There's currently no teardown done for Redis.
+        We delete the keys in redis for tables/views being removed.
         """
-        pass
+        for table in tables:
+            self.delete_table_values(config, table)
 
     @staticmethod
     def _parse_connection_string(connection_string: str):
@@ -151,7 +172,6 @@ def online_write_batch(
         ex = Timestamp()
         ex.seconds = EX_SECONDS
         ex_str = ex.SerializeToString()
-
         for entity_key, values, timestamp, created_ts in data:
             redis_key_bin = _redis_key(project, entity_key)
             ts = Timestamp()
