Skip to content

Commit 0a8bf95

Browse files
authored
Fix native types multiple entities retrieval (#977)
1 parent 35a9afc commit 0a8bf95

2 files changed

Lines changed: 63 additions & 7 deletions

File tree

sdk/python/feast/client.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -993,6 +993,7 @@ def _infer_online_entity_rows(
993993
entity_type_map = dict()
994994

995995
for entity in entity_rows_dicts:
996+
fields = {}
996997
for key, value in entity.items():
997998
# Allow for feast.types.Value
998999
if isinstance(value, Value):
@@ -1009,9 +1010,8 @@ def _infer_online_entity_rows(
10091010
f"Input entity {key} has mixed types, {current_dtype} and {entity_type_map[key]}. That is not allowed. "
10101011
)
10111012
proto_value = _python_value_to_proto_value(current_dtype, value)
1012-
entity_row_list.append(
1013-
GetOnlineFeaturesRequest.EntityRow(fields={key: proto_value})
1014-
)
1013+
fields[key] = proto_value
1014+
entity_row_list.append(GetOnlineFeaturesRequest.EntityRow(fields=fields))
10151015
return entity_row_list
10161016

10171017

tests/e2e/redis/basic-ingest-redis-serving.py

Lines changed: 60 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -709,6 +709,62 @@ def try_get_features():
709709

710710
@pytest.mark.timeout(600)
711711
@pytest.mark.run(order=18)
712+
def test_basic_ingest_retrieval_multi_entities(client):
713+
# Set to another project to test ingestion based on current project context
714+
client.set_project(PROJECT_NAME + "_NS1")
715+
merchant_fs = FeatureSet(
716+
name="merchant_fs",
717+
features=[Feature(name="merchant_sales", dtype=ValueType.FLOAT)],
718+
entities=[
719+
Entity("driver_id", ValueType.INT64),
720+
Entity("merchant_id", ValueType.INT64),
721+
],
722+
max_age=Duration(seconds=3600),
723+
)
724+
client.apply(merchant_fs)
725+
726+
N_ROWS = 2
727+
time_offset = datetime.utcnow().replace(tzinfo=pytz.utc)
728+
merchant_df = pd.DataFrame(
729+
{
730+
"datetime": [time_offset] * N_ROWS,
731+
"driver_id": [i for i in range(N_ROWS)],
732+
"merchant_id": [i for i in range(N_ROWS)],
733+
"merchant_sales": [float(i) + 0.5 for i in range(N_ROWS)],
734+
}
735+
)
736+
client.ingest("merchant_fs", merchant_df, timeout=600)
737+
738+
online_request_entity = [
739+
{"driver_id": 0, "merchant_id": 0},
740+
{"driver_id": 1, "merchant_id": 1},
741+
]
742+
online_request_features = ["merchant_sales"]
743+
744+
def try_get_features():
745+
response = client.get_online_features(
746+
entity_rows=online_request_entity, feature_refs=online_request_features
747+
)
748+
is_ok = check_online_response("merchant_sales", merchant_df, response)
749+
return response, is_ok
750+
751+
online_features_actual = wait_retry_backoff(
752+
retry_fn=try_get_features,
753+
timeout_secs=90,
754+
timeout_msg="Timed out trying to get online feature values",
755+
)
756+
757+
online_features_expected = {
758+
"driver_id": [0, 1],
759+
"merchant_id": [0, 1],
760+
"merchant_sales": [0.5, 1.5],
761+
}
762+
763+
assert online_features_actual.to_dict() == online_features_expected
764+
765+
766+
@pytest.mark.timeout(600)
767+
@pytest.mark.run(order=19)
712768
def test_basic_retrieve_feature_row_missing_fields(client, cust_trans_df):
713769
feature_refs = ["daily_transactions", "total_transactions", "null_values"]
714770

@@ -756,7 +812,7 @@ def try_get_features():
756812

757813

758814
@pytest.mark.timeout(600)
759-
@pytest.mark.run(order=19)
815+
@pytest.mark.run(order=20)
760816
def test_basic_retrieve_feature_row_extra_fields(client, cust_trans_df):
761817
feature_refs = ["daily_transactions", "total_transactions"]
762818
# apply cust_trans_fs and ingest dataframe
@@ -851,7 +907,7 @@ def all_types_dataframe():
851907

852908

853909
@pytest.mark.timeout(45)
854-
@pytest.mark.run(order=20)
910+
@pytest.mark.run(order=21)
855911
def test_all_types_register_feature_set_success(client):
856912
client.set_project(PROJECT_NAME)
857913

@@ -897,7 +953,7 @@ def test_all_types_register_feature_set_success(client):
897953

898954

899955
@pytest.mark.timeout(300)
900-
@pytest.mark.run(order=21)
956+
@pytest.mark.run(order=22)
901957
def test_all_types_ingest_success(client, all_types_dataframe):
902958
# Get all_types feature set
903959
all_types_fs = client.get_feature_set(name="all_types")
@@ -907,7 +963,7 @@ def test_all_types_ingest_success(client, all_types_dataframe):
907963

908964

909965
@pytest.mark.timeout(90)
910-
@pytest.mark.run(order=22)
966+
@pytest.mark.run(order=23)
911967
def test_all_types_retrieve_online_success(client, all_types_dataframe):
912968
# Poll serving for feature values until the correct values are returned_float_list
913969
feature_refs = [

0 commit comments

Comments
 (0)