feast-dev · HaoXuAI · Apr 15, 2024 · Mar 31, 2024 · Mar 31, 2024 · Mar 31, 2024
@@ -1690,6 +1690,67 @@ def _get_online_features(
         )
         return OnlineResponse(online_features_response)
 
+    @log_exceptions_and_usage
+    def retrieve_online_documents(
+        self,
+        feature: str,
+        query: Union[str, List[float]],
+        top_k: int,
+    ) -> OnlineResponse:
+        """
+        Retrieves the top k closest document features.
-        Retrieves the top k closest document features.
+        Retrieves the top k closest document features. Note, embeddings are a subset of features.
-        Retrieves the top k closest document features.
+        Retrieves the top k closest document features. Note, embeddings are a subset of 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".
+            query: The query to retrieve the closest document features for.
+            top_k: The number of closest document features to retrieve.
+        """
+        return self._retrieve_online_documents(
+            feature=feature,
+            query=query,
+            top_k=top_k,
+        )
+
+    def _retrieve_online_documents(
+        self,
+        feature: str,
+        query: Union[str, List[float]],
+        top_k: int,
+    ):
+        if isinstance(query, str):
+            raise ValueError(
+                "Using embedding functionality is not supported for document retrieval. Please embed the query before calling retrieve_online_documents."
+            )
+        (
+            requested_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._retrieve_from_online_store(
+            provider,
+            requested_feature_views[0],
+            requested_feature,
+            query,
+            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]]]
@@ -1906,6 +1967,48 @@ def _read_from_online_store(
             read_row_protos.append((event_timestamps, statuses, values))
         return read_row_protos
 
+    def _retrieve_from_online_store(
+        self,
+        provider: Provider,
+        table: FeatureView,
+        requested_feature: str,
+        query: List[float],
+        top_k: int,
+    ) -> List[Tuple[List[Timestamp], List["FieldStatus.ValueType"], List[Value]]]:
+        """
+        Search and return document features from the online document store.
+        """
+        documents = provider.retrieve_online_documents(
+            config=self.config,
+            table=table,
+            requested_feature=requested_feature,
+            query=query,
+            top_k=top_k,
+        )
+        # Each row is a set of features for a given entity key. We only need to convert
+        # the data to Protobuf once.
+        null_value = Value()
+        read_row_protos = []
+
+        for doc in documents:
+            row_ts_proto = Timestamp()
+            row_ts, feature_data = doc
+            # TODO (Ly): reuse whatever timestamp if row_ts is None?
+            if row_ts is not None:
+                row_ts_proto.FromDatetime(row_ts)
+            event_timestamps = [row_ts_proto]
+            if feature_data is None:
+                statuses = [FieldStatus.NOT_FOUND]
+                values = [null_value]
+            else:
+                statuses = []
+                values = []
+                for feature_name, feature_value in feature_data.items():
+                    statuses.append(FieldStatus.PRESENT)
+                    values.append(feature_value)
+            read_row_protos.append((event_timestamps, statuses, values))
+        return read_row_protos
+
     @staticmethod
     def _populate_response_from_feature_data(
         feature_data: Iterable[

@@ -21,6 +21,15 @@
 from feast.repo_config import RepoConfig
 from feast.usage import log_exceptions_and_usage
 
+# 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 feature_name, value, event_ts FROM {table_name}
+WHERE feature_name = '{feature_name}'
+ORDER BY value <-> %s
+LIMIT %s;
+"""
+
 
 class PostgreSQLOnlineStoreConfig(PostgreSQLConfig):
     type: Literal["postgres"] = "postgres"
@@ -251,6 +260,49 @@ def teardown(
             logging.exception("Teardown failed")
             raise
 
+    def retrieve_online_documents(
+        self,
+        config: RepoConfig,
+        table: FeatureView,
+        requested_feature: str,
+        embedding: List[float],
+        top_k: int,
+    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+        """
+
+        Args:
+            config: Feast configuration object
+            table: FeatureView object as the table to search
+            requested_feature: The requested feature as the column to search
+            embedding: The query embedding to search for
+            top_k: The number of items to return
+        Returns:
+            List of tuples containing the event timestamp and the document feature
+
+        """
+
+        # Convert the embedding to a string to be used in postgres vector search
+        query_embedding_str = f"'[{','.join(str(el) for el in embedding)}]'"
+
+        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
+                ),
+                (query_embedding_str, top_k),
+            )
+            rows = cur.fetchall()
+
+            for feature_name, value, event_ts in rows:
+                val = ValueProto()
+                val.ParseFromString(value)
+
+                res = {feature_name: val}
+                result.append((event_ts, res))
+
+        return result
+
 
 def _table_id(project: str, table: FeatureView) -> str:
     return f"{project}_{table.name}"

@@ -134,3 +134,29 @@ def teardown(
             entities: Entities whose corresponding infrastructure should be deleted.
         """
         pass
+
+    @abstractmethod
+    def retrieve_online_documents(
+        self,
+        config: RepoConfig,
+        table: FeatureView,
+        requested_feature: str,
+        embedding: List[float],
+        top_k: int,
+    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+        """
+        Retrieves online feature values for the specified embeddings.
+
+        Args:
+            config: The config for the current feature store.
+            table: The feature view whose feature values should be read.
+            requested_feature: The name of the feature whose embeddings should be used for retrieval.
+            embedding: The embeddings to use for retrieval.
+            top_k: The number of nearest neighbors to retrieve.
+
+        Returns:
+            object: A list of top k closest documents to the specified embedding. Each item in the list is a tuple
+            where the first item is the event timestamp for the row, and the second item is a dict of feature
+            name to embeddings.
+        """
+        pass
@@ -47,6 +47,7 @@ def __init__(self, config: RepoConfig):
         self.repo_config = config
         self._offline_store = None
         self._online_store = None
+        self._document_store = None
         self._batch_engine: Optional[BatchMaterializationEngine] = None
 
     @property
@@ -190,6 +191,23 @@ def online_read(
             )
         return result
 
+    @log_exceptions_and_usage(sampler=RatioSampler(ratio=0.001))
+    def retrieve_online_documents(
+        self,
+        config: RepoConfig,
+        table: FeatureView,
+        requested_feature: str,
+        embedding: List[float],
+        top_k: int,
+    ) -> List:
+        set_usage_attribute("provider", self.__class__.__name__)
+        result = []
+        if self.online_store:
+            result = self.online_store.retrieve_online_documents(
+                config, table, requested_feature, embedding, top_k
+            )
+        return result
+
     def ingest_df(
         self,
         feature_view: FeatureView,

@@ -295,6 +295,30 @@ def get_feature_server_endpoint(self) -> Optional[str]:
         """Returns endpoint for the feature server, if it exists."""
         return None
 
+    @abstractmethod
+    def retrieve_online_documents(
+        self,
+        config: RepoConfig,
+        table: FeatureView,
+        requested_feature: str,
+        query: List[float],
+        top_k: int,
+    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+        """
+        Searches for the top-k nearest neighbors of the given document in the online document store.
+
+        Args:
+            config: The config for the current feature store.
+            table: The feature view whose embeddings should be searched.
+            requested_feature: the requested document feature name.
+            query: The query embedding to search for.
+            top_k: The number of nearest neighbors to return.
+
+        Returns:
+            A list of dictionaries, where each dictionary contains the document feature.
+        """
+        pass
+
 
 def get_provider(config: RepoConfig) -> Provider:
     if "." not in config.provider: