Skip to content

Commit 11c00d4

Browse files
haoxu0haoxu0
authored
fix: Fix vector store config (#4583)
1 parent 1f17caa commit 11c00d4

8 files changed

Lines changed: 33 additions & 33 deletions

File tree

‎docs/reference/alpha-vector-database.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ registry:
4040
path: postgresql://@localhost:5432/feast
4141
online_store:
4242
type: postgres
43-
pgvector_enabled: true
43+
vector_enabled: true
4444
vector_len: 384
4545
host: 127.0.0.1
4646
port: 5432

‎docs/reference/online-stores/postgres.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ online_store:
3030
sslkey_path: /path/to/client-key.pem
3131
sslcert_path: /path/to/client-cert.pem
3232
sslrootcert_path: /path/to/server-ca.pem
33-
pgvector_enabled: false
33+
vector_enabled: false
3434
vector_len: 512
3535
```
3636
{% endcode %}
@@ -65,7 +65,7 @@ To compare this set of functionality against other online stores, please see the
6565
6666
## PGVector
6767
The Postgres online store supports the use of [PGVector](https://github.com/pgvector/pgvector) for storing feature values.
68-
To enable PGVector, set `pgvector_enabled: true` in the online store configuration.
68+
To enable PGVector, set `vector_enabled: true` in the online store configuration.
6969

7070
The `vector_len` parameter can be used to specify the length of the vector. The default value is 512.
7171

‎sdk/python/feast/infra/online_stores/contrib/elasticsearch.py‎

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,14 @@
1414
serialize_entity_key,
1515
)
1616
from feast.infra.online_stores.online_store import OnlineStore
17+
from feast.infra.online_stores.vector_store import VectorStoreConfig
1718
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
1819
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
1920
from feast.repo_config import FeastConfigBaseModel
2021
from feast.utils import _build_retrieve_online_document_record, to_naive_utc
2122

2223

23-
class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel):
24+
class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
2425
"""
2526
Configuration for the ElasticSearch online store.
2627
NOTE: The class *must* end with the `OnlineStoreConfig` suffix.
@@ -38,13 +39,6 @@ class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel):
3839
# The number of rows to write in a single batch
3940
write_batch_size: Optional[int] = 40
4041

41-
# The length of the vector value
42-
vector_len: Optional[int] = 512
43-
44-
# The vector similarity metric to use in KNN search
45-
# more details: https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html
46-
similarity: Optional[str] = "cosine"
47-
4842

4943
class ElasticSearchOnlineStore(OnlineStore):
5044
_client: Optional[Elasticsearch] = None

‎sdk/python/feast/infra/online_stores/contrib/postgres.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
from feast.infra.key_encoding_utils import get_list_val_str, serialize_entity_key
2626
from feast.infra.online_stores.helpers import _to_naive_utc
2727
from feast.infra.online_stores.online_store import OnlineStore
28+
from feast.infra.online_stores.vector_store import VectorStoreConfig
2829
from feast.infra.utils.postgres.connection_utils import (
2930
_get_conn,
3031
_get_conn_async,
@@ -45,15 +46,9 @@
4546
}
4647

4748

48-
class PostgreSQLOnlineStoreConfig(PostgreSQLConfig):
49+
class PostgreSQLOnlineStoreConfig(PostgreSQLConfig, VectorStoreConfig):
4950
type: Literal["postgres"] = "postgres"
5051

51-
# Whether to enable the pgvector extension for vector similarity search
52-
pgvector_enabled: Optional[bool] = False
53-
54-
# If pgvector is enabled, the length of the vector field
55-
vector_len: Optional[int] = 512
56-
5752

5853
class PostgreSQLOnlineStore(OnlineStore):
5954
_conn: Optional[Connection] = None
@@ -118,7 +113,7 @@ def online_write_batch(
118113

119114
for feature_name, val in values.items():
120115
vector_val = None
121-
if config.online_store.pgvector_enabled:
116+
if config.online_store.vector_enabled:
122117
vector_val = get_list_val_str(val)
123118
insert_values.append(
124119
(
@@ -302,7 +297,7 @@ def update(
302297

303298
for table in tables_to_keep:
304299
table_name = _table_id(project, table)
305-
if config.online_store.pgvector_enabled:
300+
if config.online_store.vector_enabled:
306301
vector_value_type = f"vector({config.online_store.vector_len})"
307302
else:
308303
# keep the vector_value_type as BYTEA if pgvector is not enabled, to maintain compatibility
@@ -380,7 +375,7 @@ def retrieve_online_documents(
380375
"""
381376
project = config.project
382377

383-
if not config.online_store.pgvector_enabled:
378+
if not config.online_store.vector_enabled:
384379
raise ValueError(
385380
"pgvector is not enabled in the online store configuration"
386381
)

‎sdk/python/feast/infra/online_stores/sqlite.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from feast.infra.infra_object import SQLITE_INFRA_OBJECT_CLASS_TYPE, InfraObject
3030
from feast.infra.key_encoding_utils import serialize_entity_key
3131
from feast.infra.online_stores.online_store import OnlineStore
32+
from feast.infra.online_stores.vector_store import VectorStoreConfig
3233
from feast.protos.feast.core.InfraObject_pb2 import InfraObject as InfraObjectProto
3334
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
3435
from feast.protos.feast.core.SqliteTable_pb2 import SqliteTable as SqliteTableProto
@@ -38,7 +39,7 @@
3839
from feast.utils import _build_retrieve_online_document_record, to_naive_utc
3940

4041

41-
class SqliteOnlineStoreConfig(FeastConfigBaseModel):
42+
class SqliteOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
4243
"""Online store config for local (SQLite-based) store"""
4344

4445
type: Literal["sqlite", "feast.infra.online_stores.sqlite.SqliteOnlineStore"] = (
@@ -49,12 +50,6 @@ class SqliteOnlineStoreConfig(FeastConfigBaseModel):
4950
path: StrictStr = "data/online.db"
5051
""" (optional) Path to sqlite db """
5152

52-
vec_enabled: Optional[bool] = False
53-
""" (optional) Enable or disable sqlite-vss for vector search"""
54-
55-
vector_len: Optional[int] = 512
56-
""" (optional) Length of the vector to be stored in the database"""
57-
5853

5954
class SqliteOnlineStore(OnlineStore):
6055
"""
@@ -83,7 +78,7 @@ def _get_conn(self, config: RepoConfig):
8378
if not self._conn:
8479
db_path = self._get_db_path(config)
8580
self._conn = _initialize_conn(db_path)
86-
if sys.version_info[0:2] == (3, 10) and config.online_store.vec_enabled:
81+
if sys.version_info[0:2] == (3, 10) and config.online_store.vector_enabled:
8782
import sqlite_vec # noqa: F401
8883

8984
self._conn.enable_load_extension(True) # type: ignore
@@ -121,7 +116,7 @@ def online_write_batch(
121116

122117
table_name = _table_id(project, table)
123118
for feature_name, val in values.items():
124-
if config.online_store.vec_enabled:
119+
if config.online_store.vector_enabled:
125120
vector_bin = serialize_f32(
126121
val.float_list_val.val, config.online_store.vector_len
127122
) # type: ignore
@@ -321,7 +316,7 @@ def retrieve_online_documents(
321316
"""
322317
project = config.project
323318

324-
if not config.online_store.vec_enabled:
319+
if not config.online_store.vector_enabled:
325320
raise ValueError("sqlite-vss is not enabled in the online store config")
326321

327322
conn = self._get_conn(config)

‎sdk/python/feast/infra/online_stores/vector_store.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
from typing import Optional
2+
3+
4+
class VectorStoreConfig:
5+
# Whether to enable the online store for vector similarity search,
6+
# This is only applicable for online store.
7+
vector_enabled: Optional[bool] = False
8+
9+
# If vector is enabled, the length of the vector field
10+
vector_len: Optional[int] = 512
11+
12+
# The vector similarity metric to use in KNN search
13+
# It is helpful for vector database that does not support config at retrieval runtime
14+
# E.g. Elasticsearch dense_vector field at
15+
# https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html
16+
similarity: Optional[str] = "cosine"

‎sdk/python/tests/integration/feature_repos/universal/online_store/postgres.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ def create_online_store(self) -> Dict[str, Any]:
6767
"user": "root",
6868
"password": "test!@#$%",
6969
"database": "test",
70-
"pgvector_enabled": True,
70+
"vector_enabled": True,
7171
"vector_len": 2,
7272
"port": self.container.get_exposed_port(5432),
7373
}

‎sdk/python/tests/unit/online_store/test_online_retrieval.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -441,7 +441,7 @@ def test_sqlite_get_online_documents() -> None:
441441
with runner.local_repo(
442442
get_example_repo("example_feature_repo_1.py"), "file"
443443
) as store:
444-
store.config.online_store.vec_enabled = True
444+
store.config.online_store.vector_enabled = True
445445
store.config.online_store.vector_len = vector_length
446446
# Write some data to two tables
447447
document_embeddings_fv = store.get_feature_view(name="document_embeddings")

0 commit comments

Comments
 (0)