Skip to content

Commit

Permalink
Refactor the datastore online_read method to be slightly more efficient
Browse files Browse the repository at this point in the history
Signed-off-by: Achal Shah <[email protected]>
  • Loading branch information
achals committed Aug 31, 2021
1 parent e895bc4 commit 97402a2
Showing 1 changed file with 9 additions and 9 deletions.
18 changes: 9 additions & 9 deletions sdk/python/feast/infra/online_stores/datastore.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down

0 comments on commit 97402a2

Please sign in to comment.