[FSTORE-1047] Support Similarity Search in the Feature Store (#1139)

* fall back to use raw feature name

cache required sk

embedding fg right join

fv get vector

fv find neighbour

write to single index

read embedding metadata

read embedding

create embedding

* check embedding in query read

* get_feature_vectors

* fix filter entry by join index

* exclude fg from serving

* fix hw login

* license

* opensearch singleton

* rename to embedding_index

* reformat

* reformat

* assign feature store to feature group

* fix serving key in helper column

* reformat

* fix test

* address comments

* minor fix

* minor fix

* read offline embedding

* read offline embedding

* fix pyspark test

* reformat

* remove hopsworks dependency

* reformat

* fix elastic url

* fix elastic url

* fix style

* fix style

Author: kennethmhc

python/hsfs/connection.py

@@ -20,6 +20,7 @@
 from requests.exceptions import ConnectionError
 
 from hsfs.decorators import connected, not_connected
+from hsfs.core.opensearch import OpenSearchClientSingleton
 from hsfs import engine, client, util, usage
 from hsfs.core import (
     feature_store_api,
@@ -276,6 +277,7 @@ def close(self):
             conn.close()
             ```
         """
+        OpenSearchClientSingleton().close()
         client.stop()
         self._feature_store_api = None
         engine.stop()

python/hsfs/constructor/query.py

@@ -138,7 +138,7 @@ def read(
         """
         if not read_options:
             read_options = {}
-
+        self._check_read_supported(online)
         sql_query, online_conn = self._prep_read(online, read_options)
 
         schema = None
@@ -179,6 +179,7 @@ def show(self, n: int, online: Optional[bool] = False):
             n: Number of rows to show.
             online: Show from online storage. Defaults to `False`.
         """
+        self._check_read_supported(online)
         read_options = {}
         sql_query, online_conn = self._prep_read(online, read_options)
 
@@ -474,6 +475,16 @@ def from_response_json(cls, json_dict):
             filter=json_decamelized.get("filter", None),
         )
 
+    def _check_read_supported(self, online):
+        if not online:
+            return
+        for fg in self.featuregroups:
+            if fg.embedding_index:
+                raise FeatureStoreException(
+                    "Reading from query containing embedding is not supported."
+                    " Use `feature_view.get_feature_vector(s) instead."
+                )
+
     @classmethod
     def _hopsworks_json(cls, json_dict):
         """

python/hsfs/core/opensearch.py

@@ -0,0 +1,70 @@
+#
+#   Copyright 2023 Logical Clocks AB
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+
+from hsfs.client.exceptions import FeatureStoreException
+import logging
+from hsfs.core.opensearch_api import OpenSearchApi
+from hsfs import client
+
+
+class OpenSearchClientSingleton:
+    _instance = None
+
+    def __new__(cls):
+        if not cls._instance:
+            cls._instance = super(OpenSearchClientSingleton, cls).__new__(cls)
+            cls._instance._opensearch_client = None
+            cls._instance._setup_opensearch_client()
+        return cls._instance
+
+    def _setup_opensearch_client(self):
+        if not self._opensearch_client:
+            try:
+                from opensearchpy import OpenSearch
+                from opensearchpy.exceptions import (
+                    ConnectionError as OpenSearchConnectionError,
+                )
+
+                self.OpenSearchConnectionError = OpenSearchConnectionError
+            except ModuleNotFoundError:
+                raise FeatureStoreException(
+                    "hopsworks and opensearchpy are required for embedding similarity search"
+                )
+            # query log is at INFO level
+            # 2023-11-24 15:10:49,470 INFO: POST https://localhost:9200/index/_search [status:200 request:0.041s]
+            logging.getLogger("opensearchpy").setLevel(logging.WARNING)
+            self._opensearch_client = OpenSearch(
+                **OpenSearchApi(
+                    client.get_instance()._project_id,
+                    client.get_instance()._project_name,
+                ).get_default_py_config()
+            )
+
+    def _refresh_opensearch_connection(self):
+        self._opensearch_client.close()
+        self._opensearch_client = None
+        self._setup_opensearch_client()
+
+    def search(self, index=None, body=None):
+        try:
+            return self._opensearch_client.search(body=body, index=index)
+        except self.OpenSearchConnectionError:
+            self._refresh_opensearch_connection()
+            return self._opensearch_client.search(body=body, index=index)
+
+    def close(self):
+        if self._opensearch_client:
+            self._opensearch_client.close()

python/hsfs/core/opensearch_api.py

@@ -0,0 +1,129 @@
+#
+#   Copyright 2023 Logical Clocks AB
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+
+from furl import furl
+
+from hsfs import client
+from hsfs.client.exceptions import FeatureStoreException
+from hsfs.core.variable_api import VariableApi
+
+
+class OPENSEARCH_CONFIG:
+    ELASTIC_ENDPOINT_ENV_VAR = "ELASTIC_ENDPOINT"
+    SSL_CONFIG = "es.net.ssl"
+    NODES_WAN_ONLY = "es.nodes.wan.only"
+    NODES = "es.nodes"
+    SSL_KEYSTORE_LOCATION = "es.net.ssl.keystore.location"
+    SSL_KEYSTORE_PASSWORD = "es.net.ssl.keystore.pass"
+    SSL_TRUSTSTORE_LOCATION = "es.net.ssl.truststore.location"
+    SSL_TRUSTSTORE_PASSWORD = "es.net.ssl.truststore.pass"
+    HTTP_AUTHORIZATION = "es.net.http.header.Authorization"
+    INDEX = "es.resource"
+    HOSTS = "hosts"
+    HTTP_COMPRESS = "http_compress"
+    HEADERS = "headers"
+    USE_SSL = "use_ssl"
+    VERIFY_CERTS = "verify_certs"
+    SSL_ASSERT_HOSTNAME = "ssl_assert_hostname"
+    CA_CERTS = "ca_certs"
+
+
+class OpenSearchApi:
+    def __init__(
+        self,
+        project_id,
+        project_name,
+    ):
+        self._project_id = project_id
+        self._project_name = project_name
+        self._variable_api = VariableApi()
+
+    def _get_opensearch_url(self):
+        if isinstance(client.get_instance(), client.external.Client):
+            external_domain = self._variable_api.get_loadbalancer_external_domain()
+            if external_domain == "":
+                raise FeatureStoreException(
+                    "External client could not locate loadbalancer_external_domain "
+                    "in cluster configuration or variable is empty."
+                )
+            return f"https://{external_domain}:9200"
+        else:
+            service_discovery_domain = self._variable_api.get_service_discovery_domain()
+            if service_discovery_domain == "":
+                raise FeatureStoreException(
+                    "Client could not locate service_discovery_domain "
+                    "in cluster configuration or variable is empty."
+                )
+            return f"https://rest.elastic.service.{service_discovery_domain}:9200"
+
+    def get_project_index(self, index):
+        """
+        This helper method prefixes the supplied index name with the project name to avoid index name clashes.
+
+        Args:
+            :index: the opensearch index to interact with.
+
+        Returns:
+            A valid opensearch index name.
+        """
+        return (self._project_name + "_" + index).lower()
+
+    def get_default_py_config(self):
+        """
+        Get the required opensearch configuration to setup a connection using the *opensearch-py* library.
+
+        ```python
+
+        import hopsworks
+        from opensearchpy import OpenSearch
+
+        project = hopsworks.login()
+
+        opensearch_api = project.get_opensearch_api()
+
+        client = OpenSearch(**opensearch_api.get_default_py_config())
+
+        ```
+        Returns:
+            A dictionary with required configuration.
+        """
+        url = furl(self._get_opensearch_url())
+        return {
+            OPENSEARCH_CONFIG.HOSTS: [{"host": url.host, "port": url.port}],
+            OPENSEARCH_CONFIG.HTTP_COMPRESS: False,
+            OPENSEARCH_CONFIG.HEADERS: {
+                "Authorization": self._get_authorization_token()
+            },
+            OPENSEARCH_CONFIG.USE_SSL: True,
+            OPENSEARCH_CONFIG.VERIFY_CERTS: True,
+            OPENSEARCH_CONFIG.SSL_ASSERT_HOSTNAME: False,
+            OPENSEARCH_CONFIG.CA_CERTS: client.get_instance()._get_ca_chain_path(),
+        }
+
+    def _get_authorization_token(self):
+        """Get opensearch jwt token.
+
+        # Returns
+            `str`: OpenSearch jwt token
+        # Raises
+            `RestAPIError`: If unable to get the token
+        """
+
+        _client = client.get_instance()
+        path_params = ["elastic", "jwt", self._project_id]
+
+        headers = {"content-type": "application/json"}
+        return _client._send_request("GET", path_params, headers=headers)["token"]