Added test for testing the new OpenAI api

Signed-off-by: Chaitany patel <[email protected]>

Author: patelchaitany

sdk/python/feast/feature_store.py

@@ -68,7 +68,12 @@
 )
 from feast.feast_object import FeastObject
 from feast.feature_service import FeatureService
-from feast.feature_view import DUMMY_ENTITY, DUMMY_ENTITY_NAME, FeatureView
+from feast.feature_view import (
+    DUMMY_ENTITY,
+    DUMMY_ENTITY_ID,
+    DUMMY_ENTITY_NAME,
+    FeatureView,
+)
 from feast.filter_models import ComparisonFilter, CompoundFilter, convert_dict_to_filter
 from feast.inference import (
     update_data_sources_with_inferred_event_timestamp_col,
@@ -2950,6 +2955,12 @@ def retrieve_online_documents_openai(
 
         response_dict = response.to_dict()
 
+        entity_key_names = {
+            col.name
+            for col in feature_view.entity_columns
+            if col.name != DUMMY_ENTITY_ID
+        }
+
         result_data = []
         if response_dict:
             first_key = next(iter(response_dict))
@@ -2963,14 +2974,24 @@ def retrieve_online_documents_openai(
                     val = values[i] if i < len(values) else None
                     if key == "distance":
                         score = float(val) if val is not None else 0.0
-                    else:
+                    elif key not in entity_key_names:
                         attributes[key] = val
                         if isinstance(val, str):
                             content_parts.append({"type": "text", "text": val})
 
+                if entity_key_names:
+                    key_parts = [
+                        str(response_dict[k][i])
+                        for k in sorted(entity_key_names)
+                        if k in response_dict and i < len(response_dict[k])
+                    ]
+                    file_id = f"{vector_store_id}_{'_'.join(key_parts)}"
+                else:
+                    file_id = f"{vector_store_id}_{i}"
+
                 result_data.append(
                     {
-                        "file_id": f"{vector_store_id}_{i}",
+                        "file_id": file_id,
                         "filename": vector_store_id,
                         "score": score,
                         "attributes": attributes,

sdk/python/feast/infra/online_stores/elasticsearch_online_store/elasticsearch.py

@@ -46,10 +46,56 @@ class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
 
     # The number of rows to write in a single batch
     write_batch_size: Optional[int] = 40
+    enable_openai_compatible_store: Optional[bool] = False
+
+
+logger = logging.getLogger(__name__)
+
+_NUMERIC_COMPARISON_OPS = {"gt", "gte", "lt", "lte"}
+
+
+def _filters_contain_numeric_comparison(
+    filter_obj: Union[ComparisonFilter, CompoundFilter],
+) -> bool:
+    if isinstance(filter_obj, ComparisonFilter):
+        return filter_obj.type in _NUMERIC_COMPARISON_OPS and isinstance(
+            filter_obj.value, (int, float)
+        )
+    if isinstance(filter_obj, CompoundFilter):
+        return any(_filters_contain_numeric_comparison(f) for f in filter_obj.filters)
+    return False
 
 
 class ElasticSearchOnlineStore(OnlineStore):
     _client: Optional[Elasticsearch] = None
+    _index_value_num_cache: Optional[Dict[str, bool]] = None
+
+    def _index_has_value_num(self, config: RepoConfig, index_name: str) -> bool:
+        """Check the actual ES index mapping for the value_num field.
+
+        Caches the result per index so we only hit ES once.
+        """
+        if self._index_value_num_cache is None:
+            self._index_value_num_cache = {}
+        if index_name in self._index_value_num_cache:
+            return self._index_value_num_cache[index_name]
+        try:
+            mapping = self._get_client(config).indices.get_mapping(index=index_name)
+            templates = (
+                mapping.get(index_name, {})
+                .get("mappings", {})
+                .get("dynamic_templates", [])
+            )
+            for tmpl in templates:
+                for _, tmpl_body in tmpl.items():
+                    props = tmpl_body.get("mapping", {}).get("properties", {})
+                    if "value_num" in props:
+                        self._index_value_num_cache[index_name] = True
+                        return True
+        except Exception:
+            pass
+        self._index_value_num_cache[index_name] = False
+        return False
 
     def _get_client(self, config: RepoConfig) -> Elasticsearch:
         online_store_config = config.online_store
@@ -95,6 +141,7 @@ def online_write_batch(
         progress: Optional[Callable[[int], Any]],
     ) -> None:
         insert_values = []
+        include_value_num = self._index_has_value_num(config, table.name)
         grouped_docs: dict[str, dict[str, Any]] = defaultdict(
             lambda: {
                 "features": {},
@@ -116,7 +163,7 @@ def online_write_batch(
             doc_key = f"{encoded_entity_key}_{timestamp}"
 
             for feature_name, value in values.items():
-                doc = _encode_feature_value(value)
+                doc = _encode_feature_value(value, include_value_num=include_value_num)
                 grouped_docs[doc_key]["features"][feature_name] = doc
                 grouped_docs[doc_key]["timestamp"] = timestamp
                 grouped_docs[doc_key]["created_ts"] = created_ts
@@ -210,6 +257,20 @@ def create_index(self, config: RepoConfig, table: FeatureView):
             _get_feature_view_vector_field_metadata(table), "vector_length", 512
         )
 
+        feature_properties: Dict[str, Any] = {
+            "feature_value": {"type": "binary"},
+            "value_text": {"type": "text"},
+            "vector_value": {
+                "type": "dense_vector",
+                "dims": vector_field_length,
+                "index": True,
+                "similarity": config.online_store.similarity,
+            },
+        }
+
+        if getattr(config.online_store, "enable_openai_compatible_store", False):
+            feature_properties["value_num"] = {"type": "double"}
+
         index_mapping = {
             "dynamic_templates": [
                 {
@@ -218,16 +279,7 @@ def create_index(self, config: RepoConfig, table: FeatureView):
                         "match": "*",
                         "mapping": {
                             "type": "object",
-                            "properties": {
-                                "feature_value": {"type": "binary"},
-                                "value_text": {"type": "text"},
-                                "vector_value": {
-                                    "type": "dense_vector",
-                                    "dims": vector_field_length,
-                                    "index": True,
-                                    "similarity": config.online_store.similarity,
-                                },
-                            },
+                            "properties": feature_properties,
                         },
                     }
                 }
@@ -344,6 +396,7 @@ def retrieve_online_documents(
     def _translate_filters(
         self,
         filters: Optional[Union[ComparisonFilter, CompoundFilter]],
+        has_value_num: bool = False,
     ) -> List[Dict[str, Any]]:
         """Translate filter objects into Elasticsearch Query DSL filter clauses.
 
@@ -353,62 +406,75 @@ def _translate_filters(
         """
         if filters is None:
             return []
-        return [self._translate_single_filter(filters)]
+        return [self._translate_single_filter(filters, has_value_num=has_value_num)]
 
     def _translate_single_filter(
         self,
         filter_obj: Union[ComparisonFilter, CompoundFilter],
+        has_value_num: bool = False,
     ) -> Dict[str, Any]:
         if isinstance(filter_obj, ComparisonFilter):
-            return self._translate_comparison_filter(filter_obj)
+            return self._translate_comparison_filter(
+                filter_obj, has_value_num=has_value_num
+            )
         elif isinstance(filter_obj, CompoundFilter):
-            return self._translate_compound_filter(filter_obj)
+            return self._translate_compound_filter(
+                filter_obj, has_value_num=has_value_num
+            )
         raise ValueError(f"Unknown filter type: {type(filter_obj)}")
 
     def _translate_comparison_filter(
         self,
         f: ComparisonFilter,
+        has_value_num: bool = False,
     ) -> Dict[str, Any]:
-        """Translate a ComparisonFilter to an ES Query DSL clause.
+        """Translate a ComparisonFilter to an ES Query DSL clause."""
+        is_numeric = isinstance(f.value, (int, float)) and not isinstance(f.value, bool)
+        is_numeric_list = (
+            isinstance(f.value, list)
+            and f.value
+            and isinstance(f.value[0], (int, float))
+            and not isinstance(f.value[0], bool)
+        )
 
-        Feature values in Elasticsearch are stored under
-        ``<feature_name>.value_text``, so filters target that nested path.
-        """
-        field = f"{f.key}.value_text"
+        if has_value_num and (is_numeric or is_numeric_list):
+            field = f"{f.key}.value_num"
+            fmt_val = f.value
+            fmt_list = f.value if is_numeric_list else None
+        else:
+            field = f"{f.key}.value_text"
+            fmt_val = str(f.value)
+            fmt_list = [str(v) for v in f.value] if isinstance(f.value, list) else None
 
         if f.type == "eq":
-            return {"term": {field: str(f.value)}}
+            return {"term": {field: fmt_val}}
         elif f.type == "ne":
-            return {"bool": {"must_not": [{"term": {field: str(f.value)}}]}}
-        elif f.type == "gt":
-            return {"range": {field: {"gt": f.value}}}
-        elif f.type == "gte":
-            return {"range": {field: {"gte": f.value}}}
-        elif f.type == "lt":
-            return {"range": {field: {"lt": f.value}}}
-        elif f.type == "lte":
-            return {"range": {field: {"lte": f.value}}}
+            return {"bool": {"must_not": [{"term": {field: fmt_val}}]}}
+        elif f.type in ("gt", "gte", "lt", "lte"):
+            return {"range": {field: {f.type: fmt_val}}}
         elif f.type == "in":
             if not isinstance(f.value, list):
                 raise ValueError(
                     f"'in' filter requires a list value, got {type(f.value)}"
                 )
-            return {"terms": {field: [str(v) for v in f.value]}}
+            return {"terms": {field: fmt_list}}
         elif f.type == "nin":
             if not isinstance(f.value, list):
                 raise ValueError(
                     f"'nin' filter requires a list value, got {type(f.value)}"
                 )
-            return {
-                "bool": {"must_not": [{"terms": {field: [str(v) for v in f.value]}}]}
-            }
+            return {"bool": {"must_not": [{"terms": {field: fmt_list}}]}}
         raise ValueError(f"Unsupported comparison operator: {f.type}")
 
     def _translate_compound_filter(
         self,
         f: CompoundFilter,
+        has_value_num: bool = False,
     ) -> Dict[str, Any]:
-        clauses = [self._translate_single_filter(sub) for sub in f.filters]
+        clauses = [
+            self._translate_single_filter(sub, has_value_num=has_value_num)
+            for sub in f.filters
+        ]
         if f.type == "and":
             return {"bool": {"must": clauses}}
         else:
@@ -458,7 +524,24 @@ def retrieve_online_documents_v2(
         source_fields += composite_key_name
         body["_source"] = source_fields
 
-        metadata_filters = self._translate_filters(filters)
+        has_value_num = self._index_has_value_num(config, es_index)
+
+        if (
+            filters
+            and _filters_contain_numeric_comparison(filters)
+            and not has_value_num
+        ):
+            logger.warning(
+                "Numeric comparison filters (gt, gte, lt, lte) are being used "
+                "but this index does not have a 'value_num' field. Numeric "
+                "fields are stored as text, which causes lexicographic "
+                "comparison instead of numeric comparison (e.g. '9' > '100'). "
+                "To fix this, set 'enable_openai_compatible_store: true' in "
+                "your online_store config, then teardown and re-apply your "
+                "feature store to recreate indices with the value_num field."
+            )
+
+        metadata_filters = self._translate_filters(filters, has_value_num=has_value_num)
 
         if embedding:
             similarity = (distance_metric or config.online_store.similarity).lower()
@@ -575,12 +658,15 @@ def _to_value_proto(value: Any) -> ValueProto:
     return val_proto
 
 
-def _encode_feature_value(value: ValueProto) -> Dict[str, Any]:
+def _encode_feature_value(
+    value: ValueProto,
+    include_value_num: bool = False,
+) -> Dict[str, Any]:
     """
     Encode a ValueProto into a dictionary for Elasticsearch storage.
     """
     encoded_value = base64.b64encode(value.SerializeToString()).decode("utf-8")
-    result = {"feature_value": encoded_value}
+    result: Dict[str, Any] = {"feature_value": encoded_value}
     vector_val = get_list_val_str(value)
 
     if vector_val:
@@ -591,8 +677,24 @@ def _encode_feature_value(value: ValueProto) -> Dict[str, Any]:
         result["value_text"] = value.bytes_val.decode("utf-8")
     elif value.HasField("int64_val"):
         result["value_text"] = str(value.int64_val)
+        if include_value_num:
+            result["value_num"] = value.int64_val
+    elif value.HasField("int32_val"):
+        result["value_text"] = str(value.int32_val)
+        if include_value_num:
+            result["value_num"] = value.int32_val
     elif value.HasField("double_val"):
         result["value_text"] = str(value.double_val)
+        if include_value_num:
+            result["value_num"] = value.double_val
+    elif value.HasField("float_val"):
+        result["value_text"] = str(value.float_val)
+        if include_value_num:
+            result["value_num"] = value.float_val
+    elif value.HasField("bool_val"):
+        result["value_text"] = str(value.bool_val)
+        if include_value_num:
+            result["value_num"] = 1.0 if value.bool_val else 0.0
     return result