feast-dev · franciscojavierarceo · Sep 28, 2024 · Sep 28, 2024 · Sep 28, 2024 · Sep 28, 2024
@@ -40,7 +40,7 @@ registry:
   path: postgresql://@localhost:5432/feast
 online_store:
   type: postgres
-  pgvector_enabled: true
+  vector_enabled: true
   vector_len: 384
   host: 127.0.0.1
   port: 5432

@@ -30,7 +30,7 @@ online_store:
     sslkey_path: /path/to/client-key.pem
     sslcert_path: /path/to/client-cert.pem
     sslrootcert_path: /path/to/server-ca.pem
-    pgvector_enabled: false
+    vector_enabled: false
     vector_len: 512
 ```
 {% endcode %}
@@ -65,7 +65,7 @@ To compare this set of functionality against other online stores, please see the
 
 ## PGVector
 The Postgres online store supports the use of [PGVector](https://github.com/pgvector/pgvector) for storing feature values.
-To enable PGVector, set `pgvector_enabled: true` in the online store configuration. 
+To enable PGVector, set `vector_enabled: true` in the online store configuration. 
 
 The `vector_len` parameter can be used to specify the length of the vector. The default value is 512.
 

@@ -14,13 +14,14 @@
     serialize_entity_key,
 )
 from feast.infra.online_stores.online_store import OnlineStore
+from feast.infra.online_stores.vector_store import VectorStoreConfig
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.repo_config import FeastConfigBaseModel
 from feast.utils import _build_retrieve_online_document_record, to_naive_utc
 
 
-class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel):
+class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
     """
     Configuration for the ElasticSearch online store.
     NOTE: The class *must* end with the `OnlineStoreConfig` suffix.
@@ -38,13 +39,6 @@ class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel):
     # The number of rows to write in a single batch
     write_batch_size: Optional[int] = 40
 
-    # The length of the vector value
-    vector_len: Optional[int] = 512
-
-    # The vector similarity metric to use in KNN search
-    # more details: https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html
-    similarity: Optional[str] = "cosine"
-
 
 class ElasticSearchOnlineStore(OnlineStore):
     _client: Optional[Elasticsearch] = None

@@ -25,6 +25,7 @@
 from feast.infra.key_encoding_utils import get_list_val_str, serialize_entity_key
 from feast.infra.online_stores.helpers import _to_naive_utc
 from feast.infra.online_stores.online_store import OnlineStore
+from feast.infra.online_stores.vector_store import VectorStoreConfig
 from feast.infra.utils.postgres.connection_utils import (
     _get_conn,
     _get_conn_async,
@@ -45,15 +46,9 @@
 }
 
 
-class PostgreSQLOnlineStoreConfig(PostgreSQLConfig):
+class PostgreSQLOnlineStoreConfig(PostgreSQLConfig, VectorStoreConfig):
     type: Literal["postgres"] = "postgres"
 
-    # Whether to enable the pgvector extension for vector similarity search
-    pgvector_enabled: Optional[bool] = False
-
-    # If pgvector is enabled, the length of the vector field
-    vector_len: Optional[int] = 512
-
 
 class PostgreSQLOnlineStore(OnlineStore):
     _conn: Optional[Connection] = None
@@ -118,7 +113,7 @@ def online_write_batch(
 
             for feature_name, val in values.items():
                 vector_val = None
-                if config.online_store.pgvector_enabled:
+                if config.online_store.vector_enabled:
                     vector_val = get_list_val_str(val)
                 insert_values.append(
                     (
@@ -302,7 +297,7 @@ def update(
 
             for table in tables_to_keep:
                 table_name = _table_id(project, table)
-                if config.online_store.pgvector_enabled:
+                if config.online_store.vector_enabled:
                     vector_value_type = f"vector({config.online_store.vector_len})"
                 else:
                     # keep the vector_value_type as BYTEA if pgvector is not enabled, to maintain compatibility
@@ -380,7 +375,7 @@ def retrieve_online_documents(
         """
         project = config.project
 
-        if not config.online_store.pgvector_enabled:
+        if not config.online_store.vector_enabled:
             raise ValueError(
                 "pgvector is not enabled in the online store configuration"
             )

@@ -29,6 +29,7 @@
 from feast.infra.infra_object import SQLITE_INFRA_OBJECT_CLASS_TYPE, InfraObject
 from feast.infra.key_encoding_utils import serialize_entity_key
 from feast.infra.online_stores.online_store import OnlineStore
+from feast.infra.online_stores.vector_store import VectorStoreConfig
 from feast.protos.feast.core.InfraObject_pb2 import InfraObject as InfraObjectProto
 from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
 from feast.protos.feast.core.SqliteTable_pb2 import SqliteTable as SqliteTableProto
@@ -38,7 +39,7 @@
 from feast.utils import _build_retrieve_online_document_record, to_naive_utc
 
 
-class SqliteOnlineStoreConfig(FeastConfigBaseModel):
+class SqliteOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
     """Online store config for local (SQLite-based) store"""
 
     type: Literal["sqlite", "feast.infra.online_stores.sqlite.SqliteOnlineStore"] = (
@@ -49,12 +50,6 @@ class SqliteOnlineStoreConfig(FeastConfigBaseModel):
     path: StrictStr = "data/online.db"
     """ (optional) Path to sqlite db """
 
-    vec_enabled: Optional[bool] = False
-    """ (optional) Enable or disable sqlite-vss for vector search"""
-
-    vector_len: Optional[int] = 512
-    """ (optional) Length of the vector to be stored in the database"""
-
 
 class SqliteOnlineStore(OnlineStore):
     """
@@ -83,7 +78,7 @@ def _get_conn(self, config: RepoConfig):
         if not self._conn:
             db_path = self._get_db_path(config)
             self._conn = _initialize_conn(db_path)
-            if sys.version_info[0:2] == (3, 10) and config.online_store.vec_enabled:
+            if sys.version_info[0:2] == (3, 10) and config.online_store.vector_enabled:
                 import sqlite_vec  # noqa: F401
 
                 self._conn.enable_load_extension(True)  # type: ignore
@@ -121,7 +116,7 @@ def online_write_batch(
 
                 table_name = _table_id(project, table)
                 for feature_name, val in values.items():
-                    if config.online_store.vec_enabled:
+                    if config.online_store.vector_enabled:
                         vector_bin = serialize_f32(
                             val.float_list_val.val, config.online_store.vector_len
                         )  # type: ignore
@@ -321,7 +316,7 @@ def retrieve_online_documents(
         """
         project = config.project
 
-        if not config.online_store.vec_enabled:
+        if not config.online_store.vector_enabled:
             raise ValueError("sqlite-vss is not enabled in the online store config")
 
         conn = self._get_conn(config)

@@ -0,0 +1,16 @@
+from typing import Optional
+
+
+class VectorStoreConfig:
+    # Whether to enable the online store for vector similarity search,
+    # This is only applicable for online store.
+    vector_enabled: Optional[bool] = False
+
+    # If vector is enabled, the length of the vector field
+    vector_len: Optional[int] = 512
+
+    # The vector similarity metric to use in KNN search
+    # It is helpful for vector database that does not support config at retrieval runtime
+    # E.g. Elasticsearch dense_vector field at
+    # https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html
+    similarity: Optional[str] = "cosine"
@@ -67,7 +67,7 @@ def create_online_store(self) -> Dict[str, Any]:
             "user": "root",
             "password": "test!@#$%",
             "database": "test",
-            "pgvector_enabled": True,
+            "vector_enabled": True,
             "vector_len": 2,
             "port": self.container.get_exposed_port(5432),
         }

@@ -441,7 +441,7 @@ def test_sqlite_get_online_documents() -> None:
     with runner.local_repo(
         get_example_repo("example_feature_repo_1.py"), "file"
     ) as store:
-        store.config.online_store.vec_enabled = True
+        store.config.online_store.vector_enabled = True
         store.config.online_store.vector_len = vector_length
         # Write some data to two tables
         document_embeddings_fv = store.get_feature_view(name="document_embeddings")