Skip to content

feat: Move get_online_features to OnlineStore interface #4319

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HaoXuAI merged 3 commits into from
Jul 1, 2024
Merged

feat: Move get_online_features to OnlineStore interface #4319

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7ffcaf3
move get_online_features to OnlineStore interface
tokoko Jun 27, 2024
d03f369
fix pydantic warnings
tokoko Jun 27, 2024
8f24355
run ruff format
tokoko Jun 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
185 changes: 10 additions & 175 deletions sdk/python/feast/feature_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1559,75 +1559,16 @@ def get_online_features(
... )
>>> online_response_dict = online_response.to_dict()
"""
if isinstance(entity_rows, list):
columnar: Dict[str, List[Any]] = {k: [] for k in entity_rows[0].keys()}
for entity_row in entity_rows:
for key, value in entity_row.items():
try:
columnar[key].append(value)
except KeyError as e:
raise ValueError(
"All entity_rows must have the same keys."
) from e

entity_rows = columnar
provider = self._get_provider()

(
join_key_values,
grouped_refs,
entity_name_to_join_key_map,
requested_on_demand_feature_views,
feature_refs,
requested_result_row_names,
online_features_response,
) = utils._prepare_entities_to_read_from_online_store(
return provider.get_online_features(
config=self.config,
features=features,
entity_rows=entity_rows,
registry=self._registry,
project=self.project,
features=features,
entity_values=entity_rows,
full_feature_names=full_feature_names,
native_entity_values=True,
)

provider = self._get_provider()
for table, requested_features in grouped_refs:
# Get the correct set of entity values with the correct join keys.
table_entity_values, idxs = utils._get_unique_entities(
table,
join_key_values,
entity_name_to_join_key_map,
)

# Fetch feature data for the minimum set of Entities.
feature_data = self._read_from_online_store(
table_entity_values,
provider,
requested_features,
table,
)

# Populate the result_rows with the Features from the OnlineStore inplace.
utils._populate_response_from_feature_data(
feature_data,
idxs,
online_features_response,
full_feature_names,
requested_features,
table,
)

if requested_on_demand_feature_views:
utils._augment_response_with_on_demand_transforms(
online_features_response,
feature_refs,
requested_on_demand_feature_views,
full_feature_names,
)

utils._drop_unneeded_columns(
online_features_response, requested_result_row_names
)
return OnlineResponse(online_features_response)

async def get_online_features_async(
self,
Expand Down Expand Up @@ -1664,75 +1605,16 @@ async def get_online_features_async(
Raises:
Exception: No entity with the specified name exists.
"""
if isinstance(entity_rows, list):
columnar: Dict[str, List[Any]] = {k: [] for k in entity_rows[0].keys()}
for entity_row in entity_rows:
for key, value in entity_row.items():
try:
columnar[key].append(value)
except KeyError as e:
raise ValueError(
"All entity_rows must have the same keys."
) from e

entity_rows = columnar
provider = self._get_provider()

(
join_key_values,
grouped_refs,
entity_name_to_join_key_map,
requested_on_demand_feature_views,
feature_refs,
requested_result_row_names,
online_features_response,
) = utils._prepare_entities_to_read_from_online_store(
return await provider.get_online_features_async(
config=self.config,
features=features,
entity_rows=entity_rows,
registry=self._registry,
project=self.project,
features=features,
entity_values=entity_rows,
full_feature_names=full_feature_names,
native_entity_values=True,
)

provider = self._get_provider()
for table, requested_features in grouped_refs:
# Get the correct set of entity values with the correct join keys.
table_entity_values, idxs = utils._get_unique_entities(
table,
join_key_values,
entity_name_to_join_key_map,
)

# Fetch feature data for the minimum set of Entities.
feature_data = await self._read_from_online_store_async(
table_entity_values,
provider,
requested_features,
table,
)

# Populate the result_rows with the Features from the OnlineStore inplace.
utils._populate_response_from_feature_data(
feature_data,
idxs,
online_features_response,
full_feature_names,
requested_features,
table,
)

if requested_on_demand_feature_views:
utils._augment_response_with_on_demand_transforms(
online_features_response,
feature_refs,
requested_on_demand_feature_views,
full_feature_names,
)

utils._drop_unneeded_columns(
online_features_response, requested_result_row_names
)
return OnlineResponse(online_features_response)

def retrieve_online_documents(
self,
Expand Down Expand Up @@ -1806,53 +1688,6 @@ def retrieve_online_documents(
)
return OnlineResponse(online_features_response)

def _read_from_online_store(
self,
entity_rows: Iterable[Mapping[str, Value]],
provider: Provider,
requested_features: List[str],
table: FeatureView,
) -> List[Tuple[List[Timestamp], List["FieldStatus.ValueType"], List[Value]]]:
"""Read and process data from the OnlineStore for a given FeatureView.

This method guarantees that the order of the data in each element of the
List returned is the same as the order of `requested_features`.

This method assumes that `provider.online_read` returns data for each
combination of Entities in `entity_rows` in the same order as they
are provided.
"""
entity_key_protos = utils._get_entity_key_protos(entity_rows)

# Fetch data for Entities.
read_rows = provider.online_read(
config=self.config,
table=table,
entity_keys=entity_key_protos,
requested_features=requested_features,
)

return utils._convert_rows_to_protobuf(requested_features, read_rows)

async def _read_from_online_store_async(
self,
entity_rows: Iterable[Mapping[str, Value]],
provider: Provider,
requested_features: List[str],
table: FeatureView,
) -> List[Tuple[List[Timestamp], List["FieldStatus.ValueType"], List[Value]]]:
entity_key_protos = utils._get_entity_key_protos(entity_rows)

# Fetch data for Entities.
read_rows = await provider.online_read_async(
config=self.config,
table=table,
entity_keys=entity_key_protos,
requested_features=requested_features,
)

return utils._convert_rows_to_protobuf(requested_features, read_rows)

def _retrieve_from_online_store(
self,
provider: Provider,
Expand Down
4 changes: 3 additions & 1 deletion sdk/python/feast/infra/materialization/kubernetes/k8s_materialization_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,9 @@ def _create_kubernetes_job(self, job_id, paths, feature_view):
def _create_configuration_map(self, job_id, paths, feature_view, namespace):
"""Create a Kubernetes configmap for this job"""

feature_store_configuration = yaml.dump(self.repo_config.dict(by_alias=True))
feature_store_configuration = yaml.dump(
self.repo_config.model_dump(by_alias=True)
)

materialization_config = yaml.dump(
{"paths": paths, "feature_view": feature_view.name}
Expand Down
Loading