Skip to content

Commit 086bed5

Browse files
patelchaitanypatelchaitany
committed
Added test for testing the new OpenAI api
Signed-off-by: Chaitany patel <patelchaitany93@gmail.com>
1 parent e96a116 commit 086bed5

8 files changed

Lines changed: 1145 additions & 328 deletions

File tree

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

Lines changed: 136 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,54 @@ class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
4646

4747
# The number of rows to write in a single batch
4848
write_batch_size: Optional[int] = 40
49+
enable_openai_compatible_store: Optional[bool] = False
50+
51+
52+
logger = logging.getLogger(__name__)
53+
54+
_NUMERIC_COMPARISON_OPS = {"gt", "gte", "lt", "lte"}
55+
56+
57+
def _filters_contain_numeric_comparison(
58+
filter_obj: Union[ComparisonFilter, CompoundFilter],
59+
) -> bool:
60+
if isinstance(filter_obj, ComparisonFilter):
61+
return filter_obj.type in _NUMERIC_COMPARISON_OPS and isinstance(
62+
filter_obj.value, (int, float)
63+
)
64+
if isinstance(filter_obj, CompoundFilter):
65+
return any(_filters_contain_numeric_comparison(f) for f in filter_obj.filters)
66+
return False
4967

5068

5169
class ElasticSearchOnlineStore(OnlineStore):
5270
_client: Optional[Elasticsearch] = None
71+
_index_value_num_cache: Dict[str, bool] = {}
72+
73+
def _index_has_value_num(self, config: RepoConfig, index_name: str) -> bool:
74+
"""Check the actual ES index mapping for the value_num field.
75+
76+
Caches the result per index so we only hit ES once.
77+
"""
78+
if index_name in self._index_value_num_cache:
79+
return self._index_value_num_cache[index_name]
80+
try:
81+
mapping = self._get_client(config).indices.get_mapping(index=index_name)
82+
templates = (
83+
mapping.get(index_name, {})
84+
.get("mappings", {})
85+
.get("dynamic_templates", [])
86+
)
87+
for tmpl in templates:
88+
for _, tmpl_body in tmpl.items():
89+
props = tmpl_body.get("mapping", {}).get("properties", {})
90+
if "value_num" in props:
91+
self._index_value_num_cache[index_name] = True
92+
return True
93+
except Exception:
94+
pass
95+
self._index_value_num_cache[index_name] = False
96+
return False
5397

5498
def _get_client(self, config: RepoConfig) -> Elasticsearch:
5599
online_store_config = config.online_store
@@ -95,6 +139,7 @@ def online_write_batch(
95139
progress: Optional[Callable[[int], Any]],
96140
) -> None:
97141
insert_values = []
142+
include_value_num = self._index_has_value_num(config, table.name)
98143
grouped_docs: dict[str, dict[str, Any]] = defaultdict(
99144
lambda: {
100145
"features": {},
@@ -116,7 +161,7 @@ def online_write_batch(
116161
doc_key = f"{encoded_entity_key}_{timestamp}"
117162

118163
for feature_name, value in values.items():
119-
doc = _encode_feature_value(value)
164+
doc = _encode_feature_value(value, include_value_num=include_value_num)
120165
grouped_docs[doc_key]["features"][feature_name] = doc
121166
grouped_docs[doc_key]["timestamp"] = timestamp
122167
grouped_docs[doc_key]["created_ts"] = created_ts
@@ -210,6 +255,20 @@ def create_index(self, config: RepoConfig, table: FeatureView):
210255
_get_feature_view_vector_field_metadata(table), "vector_length", 512
211256
)
212257

258+
feature_properties: Dict[str, Any] = {
259+
"feature_value": {"type": "binary"},
260+
"value_text": {"type": "text"},
261+
"vector_value": {
262+
"type": "dense_vector",
263+
"dims": vector_field_length,
264+
"index": True,
265+
"similarity": config.online_store.similarity,
266+
},
267+
}
268+
269+
if getattr(config.online_store, "enable_openai_compatible_store", False):
270+
feature_properties["value_num"] = {"type": "double"}
271+
213272
index_mapping = {
214273
"dynamic_templates": [
215274
{
@@ -218,16 +277,7 @@ def create_index(self, config: RepoConfig, table: FeatureView):
218277
"match": "*",
219278
"mapping": {
220279
"type": "object",
221-
"properties": {
222-
"feature_value": {"type": "binary"},
223-
"value_text": {"type": "text"},
224-
"vector_value": {
225-
"type": "dense_vector",
226-
"dims": vector_field_length,
227-
"index": True,
228-
"similarity": config.online_store.similarity,
229-
},
230-
},
280+
"properties": feature_properties,
231281
},
232282
}
233283
}
@@ -344,6 +394,7 @@ def retrieve_online_documents(
344394
def _translate_filters(
345395
self,
346396
filters: Optional[Union[ComparisonFilter, CompoundFilter]],
397+
has_value_num: bool = False,
347398
) -> List[Dict[str, Any]]:
348399
"""Translate filter objects into Elasticsearch Query DSL filter clauses.
349400
@@ -353,62 +404,74 @@ def _translate_filters(
353404
"""
354405
if filters is None:
355406
return []
356-
return [self._translate_single_filter(filters)]
407+
return [self._translate_single_filter(filters, has_value_num=has_value_num)]
357408

358409
def _translate_single_filter(
359410
self,
360411
filter_obj: Union[ComparisonFilter, CompoundFilter],
412+
has_value_num: bool = False,
361413
) -> Dict[str, Any]:
362414
if isinstance(filter_obj, ComparisonFilter):
363-
return self._translate_comparison_filter(filter_obj)
415+
return self._translate_comparison_filter(
416+
filter_obj, has_value_num=has_value_num
417+
)
364418
elif isinstance(filter_obj, CompoundFilter):
365-
return self._translate_compound_filter(filter_obj)
419+
return self._translate_compound_filter(
420+
filter_obj, has_value_num=has_value_num
421+
)
366422
raise ValueError(f"Unknown filter type: {type(filter_obj)}")
367423

368424
def _translate_comparison_filter(
369425
self,
370426
f: ComparisonFilter,
427+
has_value_num: bool = False,
371428
) -> Dict[str, Any]:
372-
"""Translate a ComparisonFilter to an ES Query DSL clause.
429+
"""Translate a ComparisonFilter to an ES Query DSL clause."""
430+
is_numeric = isinstance(f.value, (int, float)) and not isinstance(f.value, bool)
431+
is_numeric_list = (
432+
isinstance(f.value, list)
433+
and f.value
434+
and isinstance(f.value[0], (int, float))
435+
)
373436

374-
Feature values in Elasticsearch are stored under
375-
``<feature_name>.value_text``, so filters target that nested path.
376-
"""
377-
field = f"{f.key}.value_text"
437+
if has_value_num and (is_numeric or is_numeric_list):
438+
field = f"{f.key}.value_num"
439+
fmt_val = f.value
440+
fmt_list = f.value if is_numeric_list else None
441+
else:
442+
field = f"{f.key}.value_text"
443+
fmt_val = str(f.value)
444+
fmt_list = [str(v) for v in f.value] if isinstance(f.value, list) else None
378445

379446
if f.type == "eq":
380-
return {"term": {field: str(f.value)}}
447+
return {"term": {field: fmt_val}}
381448
elif f.type == "ne":
382-
return {"bool": {"must_not": [{"term": {field: str(f.value)}}]}}
383-
elif f.type == "gt":
384-
return {"range": {field: {"gt": f.value}}}
385-
elif f.type == "gte":
386-
return {"range": {field: {"gte": f.value}}}
387-
elif f.type == "lt":
388-
return {"range": {field: {"lt": f.value}}}
389-
elif f.type == "lte":
390-
return {"range": {field: {"lte": f.value}}}
449+
return {"bool": {"must_not": [{"term": {field: fmt_val}}]}}
450+
elif f.type in ("gt", "gte", "lt", "lte"):
451+
return {"range": {field: {f.type: fmt_val}}}
391452
elif f.type == "in":
392453
if not isinstance(f.value, list):
393454
raise ValueError(
394455
f"'in' filter requires a list value, got {type(f.value)}"
395456
)
396-
return {"terms": {field: [str(v) for v in f.value]}}
457+
return {"terms": {field: fmt_list}}
397458
elif f.type == "nin":
398459
if not isinstance(f.value, list):
399460
raise ValueError(
400461
f"'nin' filter requires a list value, got {type(f.value)}"
401462
)
402-
return {
403-
"bool": {"must_not": [{"terms": {field: [str(v) for v in f.value]}}]}
404-
}
463+
return {"bool": {"must_not": [{"terms": {field: fmt_list}}]}}
405464
raise ValueError(f"Unsupported comparison operator: {f.type}")
406465

407466
def _translate_compound_filter(
408467
self,
409468
f: CompoundFilter,
469+
has_value_num: bool = False,
410470
) -> Dict[str, Any]:
411-
clauses = [self._translate_single_filter(sub) for sub in f.filters]
471+
clauses = [
472+
self._translate_single_filter(sub, has_value_num=has_value_num)
473+
for sub in f.filters
474+
]
412475
if f.type == "and":
413476
return {"bool": {"must": clauses}}
414477
else:
@@ -458,7 +521,24 @@ def retrieve_online_documents_v2(
458521
source_fields += composite_key_name
459522
body["_source"] = source_fields
460523

461-
metadata_filters = self._translate_filters(filters)
524+
has_value_num = self._index_has_value_num(config, es_index)
525+
526+
if (
527+
filters
528+
and _filters_contain_numeric_comparison(filters)
529+
and not has_value_num
530+
):
531+
logger.warning(
532+
"Numeric comparison filters (gt, gte, lt, lte) are being used "
533+
"but this index does not have a 'value_num' field. Numeric "
534+
"fields are stored as text, which causes lexicographic "
535+
"comparison instead of numeric comparison (e.g. '9' > '100'). "
536+
"To fix this, set 'enable_openai_compatible_store: true' in "
537+
"your online_store config, then teardown and re-apply your "
538+
"feature store to recreate indices with the value_num field."
539+
)
540+
541+
metadata_filters = self._translate_filters(filters, has_value_num=has_value_num)
462542

463543
if embedding:
464544
similarity = (distance_metric or config.online_store.similarity).lower()
@@ -575,12 +655,15 @@ def _to_value_proto(value: Any) -> ValueProto:
575655
return val_proto
576656

577657

578-
def _encode_feature_value(value: ValueProto) -> Dict[str, Any]:
658+
def _encode_feature_value(
659+
value: ValueProto,
660+
include_value_num: bool = False,
661+
) -> Dict[str, Any]:
579662
"""
580663
Encode a ValueProto into a dictionary for Elasticsearch storage.
581664
"""
582665
encoded_value = base64.b64encode(value.SerializeToString()).decode("utf-8")
583-
result = {"feature_value": encoded_value}
666+
result: Dict[str, Any] = {"feature_value": encoded_value}
584667
vector_val = get_list_val_str(value)
585668

586669
if vector_val:
@@ -591,8 +674,24 @@ def _encode_feature_value(value: ValueProto) -> Dict[str, Any]:
591674
result["value_text"] = value.bytes_val.decode("utf-8")
592675
elif value.HasField("int64_val"):
593676
result["value_text"] = str(value.int64_val)
677+
if include_value_num:
678+
result["value_num"] = value.int64_val
679+
elif value.HasField("int32_val"):
680+
result["value_text"] = str(value.int32_val)
681+
if include_value_num:
682+
result["value_num"] = value.int32_val
594683
elif value.HasField("double_val"):
595684
result["value_text"] = str(value.double_val)
685+
if include_value_num:
686+
result["value_num"] = value.double_val
687+
elif value.HasField("float_val"):
688+
result["value_text"] = str(value.float_val)
689+
if include_value_num:
690+
result["value_num"] = value.float_val
691+
elif value.HasField("bool_val"):
692+
result["value_text"] = str(value.bool_val)
693+
if include_value_num:
694+
result["value_num"] = 1.0 if value.bool_val else 0.0
596695
return result
597696

598697

0 commit comments

Comments
 (0)