Skip to content

Commit abdb5aa

Browse files
committed
Add store validation and comments to batch retrieval
1 parent 9b2c77a commit abdb5aa

2 files changed

Lines changed: 15 additions & 1 deletion

File tree

sdk/python/feast/client.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
DatasetSource,
4444
DataFormat,
4545
FeatureSetRequest,
46+
FeastServingType,
4647
)
4748
from feast.serving.ServingService_pb2_grpc import ServingServiceStub
4849
from feast.serving.ServingService_pb2 import GetFeastServingInfoResponse
@@ -326,12 +327,21 @@ def get_batch_features(
326327

327328
try:
328329
fs_request = _build_feature_set_request(feature_ids)
330+
331+
# Validate entity rows based on entities in Feast Core
329332
self._validate_entity_rows_for_batch_retrieval(entity_rows, fs_request)
330333

334+
# Retrieve serving information to determine store type and staging location
331335
serving_info = (
332336
self._serving_service_stub.GetFeastServingInfo()
333337
) # type: GetFeastServingInfoResponse
334338

339+
if serving_info.type != FeastServingType.FEAST_SERVING_TYPE_BATCH:
340+
raise Exception(
341+
f'You are connected to a store "{self._serving_url}" which does not support batch retrieval'
342+
)
343+
344+
# Export and upload entity row dataframe to staging location provided by Feast
335345
staged_file = export_dataframe_to_staging_location(
336346
entity_rows, serving_info.job_staging_location
337347
) # type: str
@@ -344,6 +354,8 @@ def get_batch_features(
344354
)
345355
),
346356
)
357+
358+
# Retrieve Feast Job object to manage life cycle of retrieval
347359
response = self._serving_service_stub.GetBatchFeatures(request)
348360
return Job(response.job, self._serving_service_stub)
349361

sdk/python/tests/test_client.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
JobStatus,
3636
DataFormat,
3737
GetJobResponse,
38+
FeastServingType,
3839
)
3940
import pytest
4041
from feast.client import Client
@@ -292,7 +293,8 @@ def test_get_batch_features(self, mock_client, mocker):
292293
mock_client._serving_service_stub,
293294
"GetFeastServingInfo",
294295
return_value=GetFeastServingInfoResponse(
295-
job_staging_location=f"file://{tempfile.mkdtemp()}"
296+
job_staging_location=f"file://{tempfile.mkdtemp()}",
297+
type=FeastServingType.FEAST_SERVING_TYPE_BATCH,
296298
),
297299
)
298300

0 commit comments

Comments
 (0)