Updating retrieve online documents v2 to work for other fields for sqlite

franciscojavierarceo · franciscojavierarceo · commit 65cd347bbf8b · 2025-02-26T13:35:15.000-05:00
Signed-off-by: Francisco Javier Arceo &lt;farceo@redhat.com&gt;
diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py
@@ -191,6 +191,10 @@ def __init__(
             else:
                 features.append(field)
 
+        assert len([f for f in features if f.vector_index]) < 2, (
+            f"Only one vector feature is allowed per feature view. Please update {self.name}."
+        )
+
         # TODO(felixwang9817): Add more robust validation of features.
         cols = [field.name for field in schema]
         for col in cols:
diff --git a/sdk/python/feast/infra/online_stores/sqlite.py b/sdk/python/feast/infra/online_stores/sqlite.py
@@ -18,12 +18,13 @@
 import sys
 from datetime import date, datetime
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
 
 from pydantic import StrictStr
 
 from feast import Entity
 from feast.feature_view import FeatureView
+from feast.field import Field
 from feast.infra.infra_object import SQLITE_INFRA_OBJECT_CLASS_TYPE, InfraObject
 from feast.infra.key_encoding_utils import (
     deserialize_entity_key,
@@ -38,7 +39,13 @@
 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, RepoConfig
-from feast.utils import _build_retrieve_online_document_record, to_naive_utc
+from feast.type_map import feast_value_type_to_python_type
+from feast.types import FEAST_VECTOR_TYPES
+from feast.utils import (
+    _build_retrieve_online_document_record,
+    _serialize_vector_to_float_list,
+    to_naive_utc,
+)
 
 
 def adapt_date_iso(val: date):
@@ -94,6 +101,7 @@ class SqliteOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
 
     vector_enabled: bool = False
     vector_len: Optional[int] = None
+    text_search_enabled: bool = False
 
 
 class SqliteOnlineStore(OnlineStore):
@@ -144,9 +152,8 @@ def online_write_batch(
         progress: Optional[Callable[[int], Any]],
     ) -> None:
         conn = self._get_conn(config)
-
         project = config.project
-
+        feature_type_dict = {f.name: f.dtype for f in table.features}
         with conn:
             for entity_key, values, timestamp, created_ts in data:
                 entity_key_bin = serialize_entity_key(
@@ -160,71 +167,51 @@ def online_write_batch(
                 table_name = _table_id(project, table)
                 for feature_name, val in values.items():
                     if config.online_store.vector_enabled:
-                        vector_bin = serialize_f32(
-                            val.float_list_val.val, config.online_store.vector_len
-                        )  # type: ignore
-                        conn.execute(
-                            f"""
-                                    UPDATE {table_name}
-                                    SET value = ?, vector_value = ?, event_ts = ?, created_ts = ?
-                                    WHERE (entity_key = ? AND feature_name = ?)
-                                """,
-                            (
-                                # SET
-                                val.SerializeToString(),
-                                vector_bin,
-                                timestamp,
-                                created_ts,
-                                # WHERE
-                                entity_key_bin,
-                                feature_name,
-                            ),
-                        )
+                        if feature_type_dict[feature_name] in FEAST_VECTOR_TYPES:
+                            val_bin = serialize_f32(
+                                val.float_list_val.val, config.online_store.vector_len
+                            )  # type: ignore
 
+                        else:
+                            val_bin = feast_value_type_to_python_type(val)
                         conn.execute(
-                            f"""INSERT OR IGNORE INTO {table_name}
-                                (entity_key, feature_name, value, vector_value, event_ts, created_ts)
-                                VALUES (?, ?, ?, ?, ?, ?)""",
+                            f"""
+                            INSERT INTO {table_name} (entity_key, feature_name, value, vector_value, event_ts, created_ts)
+                            VALUES (?, ?, ?, ?, ?, ?)
+                            ON CONFLICT(entity_key, feature_name) DO UPDATE SET
+                                value = excluded.value,
+                                vector_value = excluded.vector_value,
+                                event_ts = excluded.event_ts,
+                                created_ts = excluded.created_ts;
+                            """,
                             (
-                                entity_key_bin,
-                                feature_name,
-                                val.SerializeToString(),
-                                vector_bin,
-                                timestamp,
-                                created_ts,
+                                entity_key_bin,  # entity_key
+                                feature_name,  # feature_name
+                                val.SerializeToString(),  # value
+                                val_bin,  # vector_value
+                                timestamp,  # event_ts
+                                created_ts,  # created_ts
                             ),
                         )
-
                     else:
                         conn.execute(
                             f"""
-                                UPDATE {table_name}
-                                SET value = ?, event_ts = ?, created_ts = ?
-                                WHERE (entity_key = ? AND feature_name = ?)
+                            INSERT INTO {table_name} (entity_key, feature_name, value, event_ts, created_ts)
+                            VALUES (?, ?, ?, ?, ?)
+                            ON CONFLICT(entity_key, feature_name) DO UPDATE SET
+                                value = excluded.value,
+                                event_ts = excluded.event_ts,
+                                created_ts = excluded.created_ts;
                             """,
                             (
-                                # SET
-                                val.SerializeToString(),
-                                timestamp,
-                                created_ts,
-                                # WHERE
-                                entity_key_bin,
-                                feature_name,
+                                entity_key_bin,  # entity_key
+                                feature_name,  # feature_name
+                                val.SerializeToString(),  # value
+                                timestamp,  # event_ts
+                                created_ts,  # created_ts
                             ),
                         )
 
-                        conn.execute(
-                            f"""INSERT OR IGNORE INTO {table_name}
-                                (entity_key, feature_name, value, event_ts, created_ts)
-                                VALUES (?, ?, ?, ?, ?)""",
-                            (
-                                entity_key_bin,
-                                feature_name,
-                                val.SerializeToString(),
-                                timestamp,
-                                created_ts,
-                            ),
-                        )
                 if progress:
                     progress(1)
 
@@ -482,13 +469,21 @@ def retrieve_online_documents_v2(
         conn = self._get_conn(config)
         cur = conn.cursor()
 
-        online_store = config.online_store
-        if not isinstance(online_store, SqliteOnlineStoreConfig):
-            raise ValueError("online_store must be SqliteOnlineStoreConfig")
         if not online_store.vector_len:
             raise ValueError("vector_len is not configured in the online store config")
+
         query_embedding_bin = serialize_f32(query, online_store.vector_len)  # type: ignore
         table_name = _table_id(config.project, table)
+        vector_fields: List[Field] = [
+            f for f in table.features if getattr(f, "vector_index", None)
+        ]
+        assert len(vector_fields) > 0, (
+            f"No vector field found, please update feature view = {table.name} to declare a vector field"
+        )
+        assert len(vector_fields) < 2, (
+            "Only one vector field is supported, please update feature view = {table.name} to declare one vector field"
+        )
+        vector_field: str = vector_fields[0].name
 
         cur.execute(
             f"""
@@ -500,17 +495,19 @@ def retrieve_online_documents_v2(
 
         cur.execute(
             f"""
-            INSERT INTO vec_table(rowid, vector_value)
+            INSERT INTO vec_table (rowid, vector_value)
             select rowid, vector_value from {table_name}
+            where feature_name = "{vector_field}"
             """
         )
 
         cur.execute(
             f"""
             select
-                fv.entity_key,
-                fv.feature_name,
-                fv.value,
+                fv2.entity_key,
+                fv2.feature_name,
+                fv2.value,
+                fv.vector_value,
                 f.distance,
                 fv.event_ts,
                 fv.created_ts
@@ -526,38 +523,80 @@ def retrieve_online_documents_v2(
             ) f
             left join {table_name} fv
                 on f.rowid = fv.rowid
-            where fv.feature_name in ({",".join(["?" for _ in requested_features])})
+            left join {table_name} fv2
+                on fv.entity_key = fv2.entity_key
+            where fv2.feature_name != "{vector_field}"
             """,
             (
                 query_embedding_bin,
                 top_k,
-                *[f.split(":")[-1] for f in requested_features],
             ),
         )
 
         rows = cur.fetchall()
-        result: List[
+        results: List[
             Tuple[
                 Optional[datetime],
                 Optional[EntityKeyProto],
                 Optional[Dict[str, ValueProto]],
             ]
         ] = []
 
-        for entity_key, feature_name, value_bin, distance, event_ts, created_ts in rows:
-            val = ValueProto()
-            val.ParseFromString(value_bin)
-            entity_key_proto = None
-            if entity_key:
-                entity_key_proto = deserialize_entity_key(
-                    entity_key,
-                    entity_key_serialization_version=config.entity_key_serialization_version,
+        entity_dict: Dict[
+            str, Dict[str, Union[str, ValueProto, EntityKeyProto, datetime]]
+        ] = {}
+        for (
+            entity_key,
+            feature_name,
+            value_bin,
+            vector_value,
+            distance,
+            event_ts,
+            created_ts,
+        ) in rows:
+            entity_key_proto = deserialize_entity_key(
+                entity_key,
+                entity_key_serialization_version=config.entity_key_serialization_version,
+            )
+            if entity_key not in entity_dict:
+                entity_dict[entity_key] = {}
+
+            feature_val = ValueProto()
+            feature_val.ParseFromString(value_bin)
+            entity_dict[entity_key]["entity_key_proto"] = entity_key_proto
+            entity_dict[entity_key][feature_name] = feature_val
+            entity_dict[entity_key][vector_field] = _serialize_vector_to_float_list(
+                vector_value
+            )
+            entity_dict[entity_key]["distance"] = ValueProto(float_val=distance)
+            entity_dict[entity_key]["event_ts"] = event_ts
+            entity_dict[entity_key]["created_ts"] = created_ts
+
+        for entity_key_value in entity_dict:
+            res_event_ts: Optional[datetime] = None
+            res_entity_key_proto: Optional[EntityKeyProto] = None
+            if isinstance(entity_dict[entity_key_value]["event_ts"], datetime):
+                res_event_ts = entity_dict[entity_key_value]["event_ts"]  # type: ignore[assignment]
+
+            if isinstance(
+                entity_dict[entity_key_value]["entity_key_proto"], EntityKeyProto
+            ):
+                res_entity_key_proto = entity_dict[entity_key_value]["entity_key_proto"]  # type: ignore[assignment]
+
+            res_dict: Dict[str, ValueProto] = {
+                k: v
+                for k, v in entity_dict[entity_key_value].items()
+                if isinstance(v, ValueProto) and isinstance(k, str)
+            }
+
+            results.append(
+                (
+                    res_event_ts,
+                    res_entity_key_proto,
+                    res_dict,
                 )
-            res = {feature_name: val}
-            res["distance"] = ValueProto(float_val=distance)
-            result.append((event_ts, entity_key_proto, res))
-
-        return result
+            )
+        return results
 
 
 def _initialize_conn(
@@ -640,7 +679,17 @@ def update(self):
             except ModuleNotFoundError:
                 logging.warning("Cannot use sqlite_vec for vector search")
         self.conn.execute(
-            f"CREATE TABLE IF NOT EXISTS {self.name} (entity_key BLOB, feature_name TEXT, value BLOB, vector_value BLOB, event_ts timestamp, created_ts timestamp,  PRIMARY KEY(entity_key, feature_name))"
+            f"""
+            CREATE TABLE IF NOT EXISTS {self.name} (
+                entity_key BLOB,
+                feature_name TEXT,
+                value BLOB,
+                vector_value BLOB,
+                event_ts timestamp,
+                created_ts timestamp,
+                PRIMARY KEY(entity_key, feature_name)
+            )
+            """
         )
         self.conn.execute(
             f"CREATE INDEX IF NOT EXISTS {self.name}_ek ON {self.name} (entity_key);"
diff --git a/sdk/python/feast/types.py b/sdk/python/feast/types.py
@@ -14,7 +14,7 @@
 from abc import ABC, abstractmethod
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Dict, Union
+from typing import Dict, List, Union
 
 import pyarrow
 
@@ -196,6 +196,17 @@ def __str__(self):
     UnixTimestamp: pyarrow.timestamp("us", tz=_utc_now().tzname()),
 }
 
+FEAST_VECTOR_TYPES: List[Union[ValueType, PrimitiveFeastType, ComplexFeastType]] = [
+    ValueType.BYTES_LIST,
+    ValueType.INT32_LIST,
+    ValueType.INT64_LIST,
+    ValueType.FLOAT_LIST,
+    ValueType.BOOL_LIST,
+]
+for k in VALUE_TYPES_TO_FEAST_TYPES:
+    if k in FEAST_VECTOR_TYPES:
+        FEAST_VECTOR_TYPES.append(VALUE_TYPES_TO_FEAST_TYPES[k])
+
 
 def from_feast_to_pyarrow_type(feast_type: FeastType) -> pyarrow.DataType:
     """
diff --git a/sdk/python/tests/example_repos/example_feature_repo_1.py b/sdk/python/tests/example_repos/example_feature_repo_1.py
@@ -125,6 +125,8 @@
             vector_search_metric="L2",
         ),
         Field(name="item_id", dtype=String),
+        Field(name="content", dtype=String),
+        Field(name="title", dtype=String),
     ],
     source=rag_documents_source,
     ttl=timedelta(hours=24),
diff --git a/sdk/python/tests/unit/online_store/test_online_retrieval.py b/sdk/python/tests/unit/online_store/test_online_retrieval.py
