def test_version(self, mock_client, mocker):
mock_client._core_service_stub = Core.CoreServiceStub(grpc.insecure_channel(""))
mock_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel("")
)
mocker.patch.object(
mock_client._core_service_stub,
"GetFeastCoreVersion",
return_value=GetFeastCoreVersionResponse(version="0.3.2"),
)
mocker.patch.object(
mock_client._serving_service_stub,
"GetFeastServingInfo",
return_value=GetFeastServingInfoResponse(version="0.3.2"),
)
status = mock_client.version()
assert (
status["core"]["url"] == CORE_URL
and status["core"]["version"] == "0.3.2"
and status["serving"]["url"] == SERVING_URL
and status["serving"]["version"] == "0.3.2"
)
def test_get_online_features(self, mocked_client, auth_metadata, mocker,
get_online_features_fields_statuses):
ROW_COUNT = 100
mocked_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
request = GetOnlineFeaturesRequestV2(project="driver_project")
request.features.extend([
FeatureRefProto(feature_table="driver", name="age"),
FeatureRefProto(feature_table="driver", name="rating"),
FeatureRefProto(feature_table="driver", name="null_value"),
])
receive_response = GetOnlineFeaturesResponse()
entity_rows = []
for row_number in range(0, ROW_COUNT):
fields = get_online_features_fields_statuses[row_number][0]
statuses = get_online_features_fields_statuses[row_number][1]
request.entity_rows.append(
GetOnlineFeaturesRequestV2.EntityRow(
fields={
"driver_id": ValueProto.Value(int64_val=row_number)
}))
entity_rows.append(
{"driver_id": ValueProto.Value(int64_val=row_number)})
receive_response.field_values.append(
GetOnlineFeaturesResponse.FieldValues(fields=fields,
statuses=statuses))
mocker.patch.object(
mocked_client._serving_service_stub,
"GetOnlineFeaturesV2",
return_value=receive_response,
)
got_response = mocked_client.get_online_features(
entity_rows=entity_rows,
feature_refs=["driver:age", "driver:rating", "driver:null_value"],
project="driver_project",
) # type: GetOnlineFeaturesResponse
mocked_client._serving_service_stub.GetOnlineFeaturesV2.assert_called_with(
request, metadata=auth_metadata, timeout=10)
got_fields = got_response.field_values[1].fields
got_statuses = got_response.field_values[1].statuses
assert (got_fields["driver_id"] == ValueProto.Value(int64_val=1)
and got_statuses["driver_id"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT
and got_fields["driver:age"] == ValueProto.Value(int64_val=1)
and got_statuses["driver:age"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT and
got_fields["driver:rating"] == ValueProto.Value(string_val="9")
and got_statuses["driver:rating"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT
and got_fields["driver:null_value"] == ValueProto.Value()
and got_statuses["driver:null_value"]
== GetOnlineFeaturesResponse.FieldStatus.NULL_VALUE)
def retrieve_job(self):
serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
job_proto = JobProto(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_RUNNING,
)
return RetrievalJob(job_proto, serving_service_stub)
def test_get_online_features(self, mock_client, mocker):
ROW_COUNT = 300
mock_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel("")
)
fields = dict()
for feature_num in range(1, 10):
fields["feature_set_1:1:feature_" + str(feature_num)] = ValueProto.Value(
int64_val=feature_num
)
field_values = GetOnlineFeaturesResponse.FieldValues(fields=fields)
response = GetOnlineFeaturesResponse()
entity_rows = []
for row_number in range(1, ROW_COUNT + 1):
response.field_values.append(field_values)
entity_rows.append(
GetOnlineFeaturesRequest.EntityRow(
fields={"customer_id": ValueProto.Value(int64_val=row_number)}
)
)
mocker.patch.object(
mock_client._serving_service_stub,
"GetOnlineFeatures",
return_value=response,
)
response = mock_client.get_online_features(
entity_rows=entity_rows,
feature_ids=[
"feature_set_1:1:feature_1",
"feature_set_1:1:feature_2",
"feature_set_1:1:feature_3",
"feature_set_1:1:feature_4",
"feature_set_1:1:feature_5",
"feature_set_1:1:feature_6",
"feature_set_1:1:feature_7",
"feature_set_1:1:feature_8",
"feature_set_1:1:feature_9",
],
) # type: GetOnlineFeaturesResponse
assert (
response.field_values[0].fields["feature_set_1:1:feature_1"].int64_val == 1
and response.field_values[0].fields["feature_set_1:1:feature_9"].int64_val
== 9
)
def test_get_online_features(self, mocked_client, mocker):
ROW_COUNT = 300
mocked_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
def int_val(x):
return ValueProto.Value(int64_val=x)
request = GetOnlineFeaturesRequest()
request.features.extend([
FeatureRefProto(project="driver_project",
feature_set="driver",
name="age"),
FeatureRefProto(project="driver_project", name="rating"),
])
recieve_response = GetOnlineFeaturesResponse()
for row_number in range(1, ROW_COUNT + 1):
request.entity_rows.append(
GetOnlineFeaturesRequest.EntityRow(
fields={"driver_id": int_val(row_number)})),
field_values = GetOnlineFeaturesResponse.FieldValues(
fields={
"driver_id": int_val(row_number),
"driver_project/driver:age": int_val(1),
"driver_project/rating": int_val(9),
})
recieve_response.field_values.append(field_values)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetOnlineFeatures",
return_value=recieve_response,
)
got_response = mocked_client.get_online_features(
entity_rows=request.entity_rows,
feature_refs=["driver:age", "rating"],
project="driver_project",
) # type: GetOnlineFeaturesResponse
mocked_client._serving_service_stub.GetOnlineFeatures.assert_called_with(
request)
got_fields = got_response.field_values[0].fields
assert (got_fields["driver_id"] == int_val(1)
and got_fields["driver:age"] == int_val(1)
and got_fields["rating"] == int_val(9))
def test_get_online_features(self, mocked_client, mocker):
ROW_COUNT = 300
mocked_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
def int_val(x):
return ValueProto.Value(int64_val=x)
# serving can return feature references with projects,
# get_online_features() should strip the project part.
field_values = GetOnlineFeaturesResponse.FieldValues(
fields={
"driver_project/driver:driver_id": int_val(1),
"driver_project/driver_id": int_val(9),
})
response = GetOnlineFeaturesResponse()
entity_rows = []
for row_number in range(1, ROW_COUNT + 1):
response.field_values.append(field_values)
entity_rows.append(
GetOnlineFeaturesRequest.EntityRow(
fields={"customer_id": int_val(row_number)}))
mocker.patch.object(
mocked_client._serving_service_stub,
"GetOnlineFeatures",
return_value=response,
)
# NOTE: Feast Serving does not allow for feature references
# that specify the same feature in the same request
response = mocked_client.get_online_features(
entity_rows=entity_rows,
feature_refs=["driver:driver_id", "driver_id"],
project="driver_project",
) # type: GetOnlineFeaturesResponse
assert (
response.field_values[0].fields["driver:driver_id"].int64_val == 1
and response.field_values[0].fields["driver_id"].int64_val == 9)
def test_get_batch_features(self, mocked_client, mocker):
mocked_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
mocked_client._core_service_stub = Core.CoreServiceStub(
grpc.insecure_channel(""))
mocker.patch.object(
mocked_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(feature_set=FeatureSetProto(
spec=FeatureSetSpecProto(
name="driver",
project="driver_project",
entities=[
EntitySpecProto(name="driver",
value_type=ValueProto.ValueType.INT64),
EntitySpecProto(
name="transaction",
value_type=ValueProto.ValueType.INT64,
),
],
features=[
FeatureSpecProto(
name="driver_id",
value_type=ValueProto.ValueType.FLOAT,
),
FeatureSpecProto(
name="driver_name",
value_type=ValueProto.ValueType.STRING,
),
],
),
meta=FeatureSetMetaProto(
status=FeatureSetStatusProto.STATUS_READY),
)),
)
expected_dataframe = pd.DataFrame({
"datetime": [datetime.utcnow() for _ in range(3)],
"driver": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
"driver_id": [1001, 1002, 1003],
})
final_results = tempfile.mktemp()
pandavro.to_avro(file_path_or_buffer=final_results,
df=expected_dataframe)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetBatchFeatures",
return_value=GetBatchFeaturesResponse(job=BatchRetrievalJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)),
)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetJob",
return_value=GetJobResponse(job=BatchRetrievalJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)),
)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetFeastServingInfo",
return_value=GetFeastServingInfoResponse(
job_staging_location=f"file://{tempfile.mkdtemp()}/",
type=FeastServingType.FEAST_SERVING_TYPE_BATCH,
),
)
mocked_client.set_project("project1")
# TODO: Abstract away GCS client and GCP dependency
# NOTE: Feast Serving does not allow for feature references
# that specify the same feature in the same request.
with patch("google.cloud.storage.Client"):
response = mocked_client.get_batch_features(
entity_rows=pd.DataFrame({
"datetime": [
pd.datetime.now(tz=timezone("Asia/Singapore"))
for _ in range(3)
],
"driver": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
}),
feature_refs=["driver:driver_id", "driver_id"],
project="driver_project",
) # Type: GetBatchFeaturesResponse
assert response.id == "123" and response.status == JobStatus.JOB_STATUS_DONE
actual_dataframe = response.to_dataframe()
assert actual_dataframe[["driver_id"
]].equals(expected_dataframe[["driver_id"]])
def test_get_batch_features(self, mock_client, mocker):
mock_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel("")
)
mock_client._core_service_stub = Core.CoreServiceStub(grpc.insecure_channel(""))
mocker.patch.object(
mock_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(
feature_set=FeatureSetSpec(
name="customer_fs",
version=1,
entities=[
EntitySpec(
name="customer", value_type=ValueProto.ValueType.INT64
),
EntitySpec(
name="transaction", value_type=ValueProto.ValueType.INT64
),
],
features=[
FeatureSpec(
name="customer_feature_1",
value_type=ValueProto.ValueType.FLOAT,
),
FeatureSpec(
name="customer_feature_2",
value_type=ValueProto.ValueType.STRING,
),
],
)
),
)
expected_dataframe = pd.DataFrame(
{
"datetime": [datetime.utcnow() for _ in range(3)],
"customer": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
"customer_fs:1:customer_feature_1": [1001, 1002, 1003],
"customer_fs:1:customer_feature_2": [1001, 1002, 1003],
}
)
final_results = tempfile.mktemp()
to_avro(file_path_or_buffer=final_results, df=expected_dataframe)
mocker.patch.object(
mock_client._serving_service_stub,
"GetBatchFeatures",
return_value=GetBatchFeaturesResponse(
job=BatchFeaturesJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)
),
)
mocker.patch.object(
mock_client._serving_service_stub,
"GetJob",
return_value=GetJobResponse(
job=BatchFeaturesJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)
),
)
mocker.patch.object(
mock_client._serving_service_stub,
"GetFeastServingInfo",
return_value=GetFeastServingInfoResponse(
job_staging_location=f"file://{tempfile.mkdtemp()}/",
type=FeastServingType.FEAST_SERVING_TYPE_BATCH,
),
)
response = mock_client.get_batch_features(
entity_rows=pd.DataFrame(
{
"datetime": [
pd.datetime.now(tz=timezone("Asia/Singapore")) for _ in range(3)
],
"customer": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
}
),
feature_ids=[
"customer_fs:1:customer_feature_1",
"customer_fs:1:customer_feature_2",
],
) # type: Job
assert response.id == "123" and response.status == JobStatus.JOB_STATUS_DONE
actual_dataframe = response.to_dataframe()
assert actual_dataframe[
["customer_fs:1:customer_feature_1", "customer_fs:1:customer_feature_2"]
].equals(
expected_dataframe[
["customer_fs:1:customer_feature_1", "customer_fs:1:customer_feature_2"]
]
)
def test_get_online_features(self, mocked_client, auth_metadata, mocker):
ROW_COUNT = 300
mocked_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
def int_val(x):
return ValueProto.Value(int64_val=x)
request = GetOnlineFeaturesRequest(project="driver_project")
request.features.extend([
FeatureRefProto(feature_set="driver", name="age"),
FeatureRefProto(name="rating"),
FeatureRefProto(name="null_value"),
])
recieve_response = GetOnlineFeaturesResponse()
entity_rows = []
for row_number in range(1, ROW_COUNT + 1):
request.entity_rows.append(
GetOnlineFeaturesRequest.EntityRow(
fields={"driver_id": int_val(row_number)}))
entity_rows.append({"driver_id": int_val(row_number)})
field_values = GetOnlineFeaturesResponse.FieldValues(
fields={
"driver_id": int_val(row_number),
"driver:age": int_val(1),
"rating": int_val(9),
"null_value": ValueProto.Value(),
},
statuses={
"driver_id": GetOnlineFeaturesResponse.FieldStatus.PRESENT,
"driver:age":
GetOnlineFeaturesResponse.FieldStatus.PRESENT,
"rating": GetOnlineFeaturesResponse.FieldStatus.PRESENT,
"null_value":
GetOnlineFeaturesResponse.FieldStatus.NULL_VALUE,
},
)
recieve_response.field_values.append(field_values)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetOnlineFeatures",
return_value=recieve_response,
)
got_response = mocked_client.get_online_features(
entity_rows=entity_rows,
feature_refs=["driver:age", "rating", "null_value"],
project="driver_project",
) # type: GetOnlineFeaturesResponse
mocked_client._serving_service_stub.GetOnlineFeatures.assert_called_with(
request, metadata=auth_metadata)
got_fields = got_response.field_values[0].fields
got_statuses = got_response.field_values[0].statuses
assert (got_fields["driver_id"] == int_val(1)
and got_statuses["driver_id"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT
and got_fields["driver:age"] == int_val(1)
and got_statuses["driver:age"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT
and got_fields["rating"] == int_val(9)
and got_statuses["rating"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT
and got_fields["null_value"] == ValueProto.Value()
and got_statuses["null_value"]
== GetOnlineFeaturesResponse.FieldStatus.NULL_VALUE)
