feat: add document store

feast-dev · HaoXuAI · Apr 15, 2024 · Mar 31, 2024 · Mar 31, 2024 · Mar 31, 2024
commit 713768e036443ecf404437685a18ee354bd087fd
@@ -35,6 +35,7 @@
 
 import pandas as pd
 import pyarrow as pa
+import numpy as np
 from colorama import Fore, Style
 from google.protobuf.timestamp_pb2 import Timestamp
 from tqdm import tqdm
@@ -1684,6 +1685,62 @@ def _get_online_features(
         )
         return OnlineResponse(online_features_response)
 
+    @log_exceptions_and_usage
+    def get_top_k_document_features(self,
+                                    feature: Union[str, FeatureService],
+                                    document: Union[str, np.ndarray],
+                                    top_k: int,
+                                    ) -> OnlineResponse:
+        """
+        Retrieves the top k cloeses document features.
-        Retrieves the top k cloeses document features.
+        Retrieves the top k closest document features.
-        Retrieves the top k cloeses document features.
+        Retrieves the top k closest document features.
+
+        Args:
+            feature: The list of document features that should be retrieved from the online document store. These features can be
+                specified either as a list of string document feature references or as a feature service. String feature
+                references must have format "feature_view:feature", e.g, "document_fv:document_embedding_feature".
-                references must have format "feature_view:feature", e.g, "document_fv:document_embedding_feature".
+                references must have format "feature_view:feature", e.g, "document_fv:document_embeddings".
-                references must have format "feature_view:feature", e.g, "document_fv:document_embedding_feature".
+                references must have format "feature_view:feature", e.g, "document_fv:document_embeddings".
+            document: The document to retrieve the closest document features for.
+            top_k: The number of closest document features to retrieve.
+        """
+        return self._get_top_k_document_features(
+            feature=feature,
+            document=document,
+            top_k=top_k,
+        )
+
+    def _get_top_k_document_features(
+            self,
+            feature: Union[str, FeatureService],
+            document: Union[str, np.ndarray],
+            top_k: int,
+    ):
+        (
+            requested_feature_views,
+            requested_on_demand_feature_views
+        ) = self._get_feature_views_to_use(
+            features=[feature],
+            allow_cache=True,
+            hide_dummy_entity=False
+        )
+        requested_feature = feature.split(":")[1] if isinstance(feature, str) else feature
+        provider = self._get_provider()
+        document_features = self._search_from_document_store(
+            provider,
+            requested_feature_views[0],
+            requested_feature,
+            document,
+            top_k,
+        )
+        online_features_response = GetOnlineFeaturesResponse(results=[])
+        self._populate_response_from_feature_data(
+            document_features,
+            [],
+            online_features_response,
+            False,
+            requested_feature,
+            requested_feature_views[0]
+        )
+        return OnlineResponse(online_features_response)
+
     @staticmethod
     def _get_columnar_entity_values(
         rowise: Optional[List[Dict[str, Any]]], columnar: Optional[Dict[str, List[Any]]]
@@ -1900,6 +1957,27 @@ def _read_from_online_store(
             read_row_protos.append((event_timestamps, statuses, values))
         return read_row_protos
 
+    def _search_from_document_store(
+        self,
+        provider: Provider,
+        table: FeatureView,
+        requested_feature: str,
+        document: Union[str, np.ndarray],
+        top_k: int,
+    ) -> List[Tuple[List[Timestamp], List["FieldStatus.ValueType"], List[Value]]]:
+        """
+        Search and return document features from the online document store.
+        """
+        documents = provider.online_search(
+            config=self.config,
+            table=table,
+            requested_feature=requested_feature,
+            document=document,
+            top_k=top_k,
+        )
+        return documents
+
+
     @staticmethod
     def _populate_response_from_feature_data(
         feature_data: Iterable[

@@ -4,6 +4,7 @@
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple
 
+import numpy as np
 import psycopg2
 import pytz
 from psycopg2 import sql
@@ -12,8 +13,10 @@
 
 from feast import Entity
 from feast.feature_view import FeatureView
+from feast.feature import Feature
 from feast.infra.key_encoding_utils import serialize_entity_key
 from feast.infra.online_stores.online_store import OnlineStore
+from feast.infra.online_stores.document_store import DocumentStore, DocumentStoreIndexConfig
 from feast.infra.utils.postgres.connection_utils import _get_conn, _get_connection_pool
 from feast.infra.utils.postgres.postgres_config import ConnectionType, PostgreSQLConfig
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
@@ -46,13 +49,13 @@ def _get_conn(self, config: RepoConfig):
 
     @log_exceptions_and_usage(online_store="postgres")
     def online_write_batch(
-        self,
-        config: RepoConfig,
-        table: FeatureView,
-        data: List[
-            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
-        ],
-        progress: Optional[Callable[[int], Any]],
+            self,
+            config: RepoConfig,
+            table: FeatureView,
+            data: List[
+                Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
+            ],
+            progress: Optional[Callable[[int], Any]],
     ) -> None:
         project = config.project
 
@@ -80,7 +83,7 @@ def online_write_batch(
             # Control the batch so that we can update the progress
             batch_size = 5000
             for i in range(0, len(insert_values), batch_size):
-                cur_batch = insert_values[i : i + batch_size]
+                cur_batch = insert_values[i: i + batch_size]
                 execute_values(
                     cur,
                     sql.SQL(
@@ -104,11 +107,11 @@ def online_write_batch(
 
     @log_exceptions_and_usage(online_store="postgres")
     def online_read(
-        self,
-        config: RepoConfig,
-        table: FeatureView,
-        entity_keys: List[EntityKeyProto],
-        requested_features: Optional[List[str]] = None,
+            self,
+            config: RepoConfig,
+            table: FeatureView,
+            entity_keys: List[EntityKeyProto],
+            requested_features: Optional[List[str]] = None,
     ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
         result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
 
@@ -175,13 +178,13 @@ def online_read(
 
     @log_exceptions_and_usage(online_store="postgres")
     def update(
-        self,
-        config: RepoConfig,
-        tables_to_delete: Sequence[FeatureView],
-        tables_to_keep: Sequence[FeatureView],
-        entities_to_delete: Sequence[Entity],
-        entities_to_keep: Sequence[Entity],
-        partial: bool,
+            self,
+            config: RepoConfig,
+            tables_to_delete: Sequence[FeatureView],
+            tables_to_keep: Sequence[FeatureView],
+            entities_to_delete: Sequence[Entity],
+            entities_to_keep: Sequence[Entity],
+            partial: bool,
     ):
         project = config.project
         schema_name = config.online_store.db_schema or config.online_store.user
@@ -236,10 +239,10 @@ def update(
             conn.commit()
 
     def teardown(
-        self,
-        config: RepoConfig,
-        tables: Sequence[FeatureView],
-        entities: Sequence[Entity],
+            self,
+            config: RepoConfig,
+            tables: Sequence[FeatureView],
+            entities: Sequence[Entity],
     ):
         project = config.project
         try:
@@ -273,3 +276,75 @@ def _to_naive_utc(ts: datetime):
         return ts
     else:
         return ts.astimezone(pytz.utc).replace(tzinfo=None)
+
+
+# Search query template to find the top k items that are closest to the given embedding
+# SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5;
+SEARCH_QUERY_TEMPLATE = """
+SELECT entity_key, feature_name, value, event_ts FROM {table_name}
+WHERE feature_name = '{feature_name}'
+ORDER BY value <-> %s
+LIMIT %s;
+"""
+
+# Create index query template to create a index based on the index type
+CREATE_INDEX_QUERY_TEMPLATE = """
+CREATE INDEX ON {table_name} USING {index_type} (embedding {embeding_type});
+"""
+
+
+class PostgresDocumentStoreConfig(DocumentStoreIndexConfig):
+    type: Literal["postgres"] = "postgres"
+
+
+class PostgresDocumentStore(PostgreSQLOnlineStore, DocumentStore):
+
+    def online_search(self,
+                      config: RepoConfig,
+                      table: FeatureView,
+                      requested_feature: str,
+                      embedding: np.ndarray,
+                      top_k: int,
+                      ):
+        result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
+
+        with self._get_conn(config) as conn, conn.cursor() as cur:
+            cur.execute(SEARCH_QUERY_TEMPLATE.format(
+                table_name=table,
+                feature_name=requested_feature
+            ), (embedding, top_k))
+            rows = cur.fetchall()
+
+            for row in rows:
+                # The first column is the entity key
+                entity_key = EntityKeyProto()
+                entity_key.ParseFromString(row[0])
+
+                # The second column is the feature name
+                feature_name = row[1]
+
+                # The third column is the embedding value
+                val = ValueProto()
+                val.ParseFromString(row[2])
+
+                # The fourth column is the event timestamp
+                event_ts = row[3]
+
+                res = {}
+                res[feature_name] = val
+                result.append((event_ts, res))
+
+
+        return result
+
+    def create_index(self,
+                     config: RepoConfig,
+                     index: str,
+                     index_config: DocumentStoreIndexConfig
+                     ):
+        with self._get_conn(config) as conn, conn.cursor() as cur:
+            cur.execute(CREATE_INDEX_QUERY_TEMPLATE.format(
+                table_name=config.project,
+                index_type=index,
+                embeding_type=index_config.embedding_type
+            ))
diff --git a/sdk/python/feast/infra/online_stores/document_store.py b/sdk/python/feast/infra/online_stores/document_store.py
@@ -0,0 +1,38 @@
+from abc import abstractmethod
+from datetime import datetime
+from feast.feature_view import FeatureView
+from feast.protos.feast.types.Value_pb2 import Value as ValueProto
+from feast.repo_config import RepoConfig, FeastConfigBaseModel
+from infra.online_stores.online_store import OnlineStore
+from typing import Optional, List, Tuple, Dict
+import numpy as np
+
+
+class DocumentStoreIndexConfig(FeastConfigBaseModel):
+    embedding_type: Optional[str]
+
+
+class DocumentStore(OnlineStore):
+    index: Optional[str]
+
+    @abstractmethod
+    def online_search(self,
+                      config: RepoConfig,
+                      table: FeatureView,
+                      requested_feature: str,
+                      embeddings: np.ndarray,
+                      top_k: int,
+                      ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+        raise NotImplementedError(
+            "You have to implement this!"
+        )
+
+    @abstractmethod
+    def create_index(self,
+                     config: RepoConfig,
+                     index: str,
+                     index_config: DocumentStoreIndexConfig
+                     ):
+        raise NotImplementedError(
+            "You have to implement this!"
+        )
diff --git a/sdk/python/feast/infra/online_stores/helpers.py b/sdk/python/feast/infra/online_stores/helpers.py
@@ -9,6 +9,7 @@
     serialize_entity_key_prefix,
 )
 from feast.infra.online_stores.online_store import OnlineStore
+from feast.infra.online_stores.document_store import DocumentStore
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 
 
@@ -21,6 +22,15 @@ def get_online_store_from_config(online_store_config: Any) -> OnlineStore:
     return online_store_class()
 
 
+def get_document_store_from_config(document_store_config: Any) -> DocumentStore:
+    """Creates a document store corresponding to the given online document store config."""
+    module_name = document_store_config.__module__
+    qualified_name = type(document_store_config).__name__
+    class_name = qualified_name.replace("Config", "")
+    document_store_class = import_class(module_name, class_name, "DocumentStore")
+    return document_store_class()
+
+
 def _redis_key(
     project: str, entity_key: EntityKeyProto, entity_key_serialization_version=1
 ) -> bytes: