From 97402a222c6313ac1f7a23253292b6590091f602 Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Tue, 31 Aug 2021 11:18:56 -0700 Subject: [PATCH 1/2] Refactor the datastore online_read method to be slightly more efficient Signed-off-by: Achal Shah --- .../feast/infra/online_stores/datastore.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/datastore.py b/sdk/python/feast/infra/online_stores/datastore.py index 767b934fd1..d2b3718314 100644 --- a/sdk/python/feast/infra/online_stores/datastore.py +++ b/sdk/python/feast/infra/online_stores/datastore.py @@ -16,6 +16,7 @@ from multiprocessing.pool import ThreadPool from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union +from google.cloud.datastore import Key from pydantic import PositiveInt, StrictStr from pydantic.typing import Literal @@ -243,22 +244,21 @@ def online_read( ) keys.append(key) + # NOTE: get_multi doesn't return values in the same order as the keys in the request. + # Also, len(values) can be less than len(keys) in the case of missing values. values = client.get_multi(keys) - - if values is not None: - keys_missing_from_response = set(keys) - set([v.key for v in values]) - values = sorted(values, key=lambda v: keys.index(v.key)) - for value in values: + values_dict = {v.key: v for v in values} if values is not None else {} + for key in keys: + if key in values_dict: + value = values_dict[key] res = {} for feature_name, value_bin in value["values"].items(): val = ValueProto() val.ParseFromString(value_bin) res[feature_name] = val result.append((value["event_ts"], res)) - for missing_key_idx in sorted( - [keys.index(k) for k in keys_missing_from_response] - ): - result.insert(missing_key_idx, (None, None)) + else: + result.append((None, None)) return result From a960eeec85830a41e91f33283bf8cd6d21ee31cf Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Tue, 31 Aug 2021 11:22:02 -0700 Subject: [PATCH 2/2] Remove double import Signed-off-by: Achal Shah --- sdk/python/feast/infra/online_stores/datastore.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sdk/python/feast/infra/online_stores/datastore.py b/sdk/python/feast/infra/online_stores/datastore.py index d2b3718314..37011c7d51 100644 --- a/sdk/python/feast/infra/online_stores/datastore.py +++ b/sdk/python/feast/infra/online_stores/datastore.py @@ -16,7 +16,6 @@ from multiprocessing.pool import ThreadPool from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union -from google.cloud.datastore import Key from pydantic import PositiveInt, StrictStr from pydantic.typing import Literal