Skip to content

Commit e844370

Browse files
franciscojavierarceofranciscojavierarceo
committed
feat: Add SQLite Offline store
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
1 parent ec6f1b7 commit e844370

3 files changed

Lines changed: 219 additions & 20 deletions

File tree

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

Lines changed: 156 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from feast.feature_view import FeatureView
2727
from feast.infra.infra_object import SQLITE_INFRA_OBJECT_CLASS_TYPE, InfraObject
2828
from feast.infra.key_encoding_utils import (
29+
deserialize_entity_key,
2930
serialize_entity_key,
3031
serialize_f32,
3132
)
@@ -91,6 +92,9 @@ class SqliteOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
9192
path: StrictStr = "data/online.db"
9293
""" (optional) Path to sqlite db """
9394

95+
vector_enabled: bool = False
96+
vector_len: Optional[int] = None
97+
9498

9599
class SqliteOnlineStore(OnlineStore):
96100
"""
@@ -104,22 +108,21 @@ class SqliteOnlineStore(OnlineStore):
104108

105109
@staticmethod
106110
def _get_db_path(config: RepoConfig) -> str:
107-
assert (
108-
config.online_store.type == "sqlite"
109-
or config.online_store.type.endswith("SqliteOnlineStore")
110-
)
111-
112-
if config.repo_path and not Path(config.online_store.path).is_absolute():
113-
db_path = str(config.repo_path / config.online_store.path)
114-
else:
115-
db_path = config.online_store.path
116-
return db_path
111+
online_store = config.online_store
112+
if not isinstance(online_store, SqliteOnlineStoreConfig):
113+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
114+
if config.repo_path and not Path(online_store.path).is_absolute():
115+
return str(config.repo_path / online_store.path)
116+
return str(online_store.path)
117117

118118
def _get_conn(self, config: RepoConfig):
119119
if not self._conn:
120120
db_path = self._get_db_path(config)
121121
self._conn = _initialize_conn(db_path)
122-
if sys.version_info[0:2] == (3, 10) and config.online_store.vector_enabled:
122+
online_store = config.online_store
123+
if not isinstance(online_store, SqliteOnlineStoreConfig):
124+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
125+
if sys.version_info[0:2] == (3, 10) and online_store.vector_enabled:
123126
import sqlite_vec # noqa: F401
124127

125128
self._conn.enable_load_extension(True) # type: ignore
@@ -141,6 +144,9 @@ def online_write_batch(
141144
],
142145
progress: Optional[Callable[[int], Any]],
143146
) -> None:
147+
online_store = config.online_store
148+
if not isinstance(online_store, SqliteOnlineStoreConfig):
149+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
144150
conn = self._get_conn(config)
145151

146152
project = config.project
@@ -157,9 +163,13 @@ def online_write_batch(
157163

158164
table_name = _table_id(project, table)
159165
for feature_name, val in values.items():
160-
if config.online_store.vector_enabled:
166+
online_store = config.online_store
167+
if not isinstance(online_store, SqliteOnlineStoreConfig):
168+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
169+
if online_store.vector_enabled and online_store.vector_len:
161170
vector_bin = serialize_f32(
162-
val.float_list_val.val, config.online_store.vector_len
171+
val.float_list_val.val,
172+
online_store.vector_len,
163173
) # type: ignore
164174
conn.execute(
165175
f"""
@@ -356,22 +366,28 @@ def retrieve_online_documents(
356366
Returns:
357367
List of tuples containing the event timestamp, the document feature, the vector value, and the distance
358368
"""
359-
project = config.project
360-
361-
if not config.online_store.vector_enabled:
369+
online_store = config.online_store
370+
if not isinstance(online_store, SqliteOnlineStoreConfig):
371+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
372+
if not online_store.vector_enabled:
362373
raise ValueError("sqlite-vss is not enabled in the online store config")
363374

364375
conn = self._get_conn(config)
365376
cur = conn.cursor()
366377

367378
# Convert the embedding to a binary format instead of using SerializeToString()
368-
query_embedding_bin = serialize_f32(embedding, config.online_store.vector_len)
369-
table_name = _table_id(project, table)
379+
online_store = config.online_store
380+
if not isinstance(online_store, SqliteOnlineStoreConfig):
381+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
382+
if not online_store.vector_len:
383+
raise ValueError("vector_len is not configured in the online store config")
384+
query_embedding_bin = serialize_f32(embedding, online_store.vector_len) # type: ignore
385+
table_name = _table_id(config.project, table)
370386

371387
cur.execute(
372388
f"""
373389
CREATE VIRTUAL TABLE vec_example using vec0(
374-
vector_value float[{config.online_store.vector_len}]
390+
vector_value float[{online_store.vector_len}]
375391
);
376392
"""
377393
)
@@ -444,6 +460,127 @@ def retrieve_online_documents(
444460

445461
return result
446462

463+
def retrieve_online_documents_v2(
464+
self,
465+
config: RepoConfig,
466+
table: FeatureView,
467+
requested_features: List[str],
468+
query: List[float],
469+
top_k: int,
470+
distance_metric: Optional[str] = None,
471+
) -> List[
472+
Tuple[
473+
Optional[datetime],
474+
Optional[EntityKeyProto],
475+
Optional[Dict[str, ValueProto]],
476+
]
477+
]:
478+
"""
479+
Retrieve documents using vector similarity search.
480+
Args:
481+
config: Feast configuration object
482+
table: FeatureView object as the table to search
483+
requested_features: List of requested features to retrieve
484+
query: Query embedding to search for
485+
top_k: Number of items to return
486+
distance_metric: Distance metric to use (optional)
487+
Returns:
488+
List of tuples containing the event timestamp, entity key, and feature values
489+
"""
490+
online_store = config.online_store
491+
if not isinstance(online_store, SqliteOnlineStoreConfig):
492+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
493+
if not online_store.vector_enabled:
494+
raise ValueError("Vector search is not enabled in the online store config")
495+
496+
conn = self._get_conn(config)
497+
cur = conn.cursor()
498+
499+
online_store = config.online_store
500+
if not isinstance(online_store, SqliteOnlineStoreConfig):
501+
raise ValueError("online_store must be SqliteOnlineStoreConfig")
502+
if not online_store.vector_len:
503+
raise ValueError("vector_len is not configured in the online store config")
504+
query_embedding_bin = serialize_f32(query, online_store.vector_len) # type: ignore
505+
table_name = _table_id(config.project, table)
506+
507+
cur.execute(
508+
f"""
509+
CREATE VIRTUAL TABLE IF NOT EXISTS vec_example using vec0(
510+
vector_value float[{online_store.vector_len}]
511+
);
512+
"""
513+
)
514+
515+
cur.execute(
516+
f"""
517+
INSERT INTO vec_example(rowid, vector_value)
518+
select rowid, vector_value from {table_name}
519+
"""
520+
)
521+
522+
cur.execute(
523+
"""
524+
INSERT INTO vec_example(rowid, vector_value)
525+
VALUES (?, ?)
526+
""",
527+
(0, query_embedding_bin),
528+
)
529+
530+
cur.execute(
531+
f"""
532+
select
533+
fv.entity_key,
534+
fv.feature_name,
535+
fv.value,
536+
f.distance,
537+
fv.event_ts,
538+
fv.created_ts
539+
from (
540+
select
541+
rowid,
542+
vector_value,
543+
distance
544+
from vec_example
545+
where vector_value match ?
546+
order by distance
547+
limit ?
548+
) f
549+
left join {table_name} fv
550+
on f.rowid = fv.rowid
551+
where fv.feature_name in ({",".join(["?" for _ in requested_features])})
552+
""",
553+
(
554+
query_embedding_bin,
555+
top_k,
556+
*[f.split(":")[-1] for f in requested_features],
557+
),
558+
)
559+
560+
rows = cur.fetchall()
561+
result: List[
562+
Tuple[
563+
Optional[datetime],
564+
Optional[EntityKeyProto],
565+
Optional[Dict[str, ValueProto]],
566+
]
567+
] = []
568+
569+
for entity_key, feature_name, value_bin, distance, event_ts, created_ts in rows:
570+
val = ValueProto()
571+
val.ParseFromString(value_bin)
572+
entity_key_proto = None
573+
if entity_key:
574+
entity_key_proto = deserialize_entity_key(
575+
entity_key,
576+
entity_key_serialization_version=config.entity_key_serialization_version,
577+
)
578+
res = {feature_name: val}
579+
res["distance"] = ValueProto(float_val=distance)
580+
result.append((event_ts, entity_key_proto, res))
581+
582+
return result
583+
447584

448585
def _initialize_conn(db_path: str):
449586
try:

‎sdk/python/tests/integration/registration/test_universal_registry.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1164,7 +1164,7 @@ def test_registry_cache_thread_async(test_registry):
11641164
test_registry.teardown()
11651165

11661166

1167-
@pytest.mark.integration
1167+
# @pytest.mark.integration
11681168
@pytest.mark.parametrize(
11691169
"test_registry",
11701170
all_fixtures,

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -738,6 +738,68 @@ def test_sqlite_vec_import() -> None:
738738
assert result == [(2, 2.39), (1, 2.39)]
739739

740740

741+
@pytest.mark.skipif(
742+
sys.version_info[0:2] != (3, 10),
743+
reason="Only works on Python 3.10",
744+
)
745+
def test_sqlite_get_online_documents_v2() -> None:
746+
"""Test retrieving documents using v2 method with vector similarity search."""
747+
n = 10
748+
vector_length = 8
749+
runner = CliRunner()
750+
with runner.local_repo(
751+
get_example_repo("example_feature_repo_1.py"), "file"
752+
) as store:
753+
store.config.online_store.vector_enabled = True
754+
store.config.online_store.vector_len = vector_length
755+
document_embeddings_fv = store.get_feature_view(name="document_embeddings")
756+
757+
provider = store._get_provider()
758+
759+
# Create test data
760+
item_keys = [
761+
EntityKeyProto(
762+
join_keys=["item_id"], entity_values=[ValueProto(int64_val=i)]
763+
)
764+
for i in range(n)
765+
]
766+
data = []
767+
for item_key in item_keys:
768+
data.append(
769+
(
770+
item_key,
771+
{
772+
"Embeddings": ValueProto(
773+
float_list_val=FloatListProto(
774+
val=[float(x) for x in np.random.random(vector_length)]
775+
)
776+
)
777+
},
778+
_utc_now(),
779+
_utc_now(),
780+
)
781+
)
782+
783+
provider.online_write_batch(
784+
config=store.config,
785+
table=document_embeddings_fv,
786+
data=data,
787+
progress=None,
788+
)
789+
790+
# Test vector similarity search
791+
query_embedding = [float(x) for x in np.random.random(vector_length)]
792+
result = store.retrieve_online_documents_v2(
793+
features=["document_embeddings:Embeddings"],
794+
query=query_embedding,
795+
top_k=3,
796+
).to_dict()
797+
798+
assert "Embeddings" in result
799+
assert "distance" in result
800+
assert len(result["distance"]) == 3
801+
802+
741803
@pytest.mark.skip(reason="Skipping this test as CI struggles with it")
742804
def test_local_milvus() -> None:
743805
import random

0 commit comments

Comments
 (0)