Skip to content

Commit 4fce072

Browse files
lclc19lclc19
committed
feat: Add dataplane code snippets for feature store service
1 parent 27a15d5 commit 4fce072

9 files changed

+521
-4
lines changed

‎samples/snippets/conftest.py‎

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from google.cloud import aiplatform, aiplatform_v1beta1
1919
from google.cloud import storage
20+
from google.cloud import bigquery
2021
import pytest
2122

2223
import helpers
@@ -89,6 +90,11 @@ def featurestore_client():
8990
)
9091
yield featurestore_client
9192

93+
@pytest.fixture
94+
def bigquery_client():
95+
bigquery_client = bigquery.Client(project=os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT"))
96+
yield bigquery_client
97+
9298
# Shared setup/teardown.
9399
@pytest.fixture()
94100
def teardown_batch_prediction_job(shared_state, job_client):
@@ -211,16 +217,24 @@ def teardown_dataset(shared_state, dataset_client):
211217
def teardown_featurestore(shared_state, featurestore_client):
212218
yield
213219

214-
# Delete the created featurestore
215-
featurestore_client.delete_featurestore(name=shared_state["featurestore_name"])
220+
# Force delete the created featurestore
221+
force_delete_featurestore_request = {
222+
"name": shared_state["featurestore_name"],
223+
"force": True
224+
}
225+
featurestore_client.delete_featurestore(request=force_delete_featurestore_request)
216226

217227

218228
@pytest.fixture()
219229
def teardown_entity_type(shared_state, featurestore_client):
220230
yield
221231

222-
# Delete the created entity type
223-
featurestore_client.delete_entity_type(name=shared_state["entity_type_name"])
232+
# Force delete the created entity type
233+
force_delete_entity_type_request = {
234+
"name": shared_state["entity_type_name"],
235+
"force": True
236+
}
237+
featurestore_client.delete_entity_type(request=force_delete_entity_type_request)
224238

225239

226240
@pytest.fixture()
@@ -231,6 +245,25 @@ def teardown_feature(shared_state, featurestore_client):
231245
featurestore_client.delete_feature(name=shared_state["feature_name"])
232246

233247

248+
@pytest.fixture()
249+
def teardown_features(shared_state, featurestore_client):
250+
yield
251+
252+
# Delete the created features
253+
for feature_name in shared_state["feature_names"]:
254+
featurestore_client.delete_feature(name=feature_name)
255+
256+
257+
@pytest.fixture()
258+
def teardown_batch_read_feature_values(shared_state, bigquery_client):
259+
yield
260+
261+
# Delete the created dataset
262+
bigquery_client.delete_dataset(
263+
shared_state["destination_data_set"], delete_contents=True, not_found_ok=True
264+
)
265+
266+
234267
@pytest.fixture()
235268
def create_endpoint(shared_state, endpoint_client):
236269
def create(project, location, test_name="temp_deploy_model_test"):

‎samples/snippets/feature_store_service/batch_create_features_sample.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START aiplatform_batch_create_features_sample]
16+
from google.cloud import aiplatform_v1beta1 as aiplatform
17+
18+
19+
def batch_create_features_sample(
20+
project: str,
21+
featurestore_id: str,
22+
entity_type_id: str,
23+
requests: str,
24+
location: str = "us-central1",
25+
api_endpoint: str = "us-central1-aiplatform.googleapis.com",
26+
timeout: int = 300,
27+
):
28+
# The AI Platform services require regional API endpoints.
29+
client_options = {"api_endpoint": api_endpoint}
30+
# Initialize client that will be used to create and send requests.
31+
# This client only needs to be created once, and can be reused for multiple requests.
32+
client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
33+
parent = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}"
34+
batch_create_features_request = {"parent": parent, "requests": requests}
35+
lro_response = client.batch_create_features(request=batch_create_features_request)
36+
print("Long running operation:", lro_response.operation.name)
37+
batch_create_features_response = lro_response.result(timeout=timeout)
38+
print("batch_create_features_response:", batch_create_features_response)
39+
40+
41+
# [END aiplatform_batch_create_features_sample]

‎samples/snippets/feature_store_service/batch_create_features_sample_test.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
from uuid import uuid4
17+
18+
import batch_create_features_sample
19+
import pytest
20+
21+
import helpers
22+
23+
PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
24+
25+
26+
@pytest.fixture(scope="function", autouse=True)
27+
def teardown(teardown_features):
28+
yield
29+
30+
31+
def test_ucaip_generated_batch_create_features_sample_vision(capsys, shared_state):
32+
featurestore_id = "perm_sample_featurestore"
33+
entity_type_id = "perm_sample_entity_type"
34+
requests = [
35+
{
36+
"feature_id": f"gender_{uuid4()}".replace("-", "_")[:60],
37+
"feature": {"value_type": "STRING", "description": "lorem ipsum"},
38+
},
39+
{
40+
"feature_id": f"liked_genres_{uuid4()}".replace("-", "_")[:60],
41+
"feature": {"value_type": "STRING_ARRAY", "description": "lorem ipsum"},
42+
},
43+
]
44+
location = "us-central1"
45+
batch_create_features_sample.batch_create_features_sample(
46+
project=PROJECT_ID,
47+
featurestore_id=featurestore_id,
48+
entity_type_id=entity_type_id,
49+
requests=requests,
50+
location=location,
51+
)
52+
out, _ = capsys.readouterr()
53+
assert "batch_create_features_response" in out
54+
55+
parent = f"projects/{PROJECT_ID}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}/features/"
56+
shared_state["feature_names"] = []
57+
for request in requests:
58+
shared_state["feature_names"].append(parent + request["feature_id"])

‎samples/snippets/feature_store_service/batch_read_feature_values_sample.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START aiplatform_batch_read_feature_values_sample]
16+
from google.cloud import aiplatform_v1beta1 as aiplatform
17+
18+
19+
def batch_read_feature_values_sample(
20+
project: str,
21+
featurestore_id: str,
22+
csv_read_instances: str,
23+
destination: str,
24+
entity_type_specs: str,
25+
location: str = "us-central1",
26+
api_endpoint: str = "us-central1-aiplatform.googleapis.com",
27+
timeout: int = 300,
28+
):
29+
# The AI Platform services require regional API endpoints.
30+
client_options = {"api_endpoint": api_endpoint}
31+
# Initialize client that will be used to create and send requests.
32+
# This client only needs to be created once, and can be reused for multiple requests.
33+
client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
34+
featurestore = (
35+
f"projects/{project}/locations/{location}/featurestores/{featurestore_id}"
36+
)
37+
batch_read_feature_values_request = {
38+
"featurestore": featurestore,
39+
"csv_read_instances": csv_read_instances,
40+
"destination": destination,
41+
"entity_type_specs": entity_type_specs,
42+
}
43+
lro_response = client.batch_read_feature_values(
44+
request=batch_read_feature_values_request
45+
)
46+
print("Long running operation:", lro_response.operation.name)
47+
batch_read_feature_values_response = lro_response.result(timeout=timeout)
48+
print("batch_read_feature_values_response:", batch_read_feature_values_response)
49+
50+
51+
# [END aiplatform_batch_read_feature_values_sample]

‎samples/snippets/feature_store_service/batch_read_feature_values_sample_test.py‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
from uuid import uuid4
17+
from datetime import datetime
18+
19+
from google.cloud import bigquery
20+
21+
import batch_read_feature_values_sample
22+
import pytest
23+
24+
import helpers
25+
26+
PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
27+
LOCATION = "us-central1"
28+
INPUT_CSV_FILE = "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/movie_prediction_perm.csv"
29+
30+
31+
@pytest.fixture(scope="function", autouse=True)
32+
def teardown(teardown_batch_read_feature_values):
33+
yield
34+
35+
36+
def setup_test():
37+
# Output dataset
38+
destination_data_set = "movie_predictions_" + datetime.now().strftime(
39+
"%Y%m%d%H%M%S"
40+
)
41+
42+
# Output table. Make sure that the table does NOT already exist; the BatchReadFeatureValues API cannot overwrite an existing table
43+
destination_table_name = "training_data"
44+
DESTINATION_PATTERN = "bq://{project}.{dataset}.{table}"
45+
destination_table_uri = DESTINATION_PATTERN.format(
46+
project=PROJECT_ID, dataset=destination_data_set, table=destination_table_name
47+
)
48+
49+
# Create dataset
50+
bq_client = bigquery.Client(project=PROJECT_ID)
51+
dataset_id = "{}.{}".format(bq_client.project, destination_data_set)
52+
dataset = bigquery.Dataset(dataset_id)
53+
dataset.location = LOCATION
54+
dataset = bq_client.create_dataset(dataset)
55+
print("Created dataset {}.{}".format(bq_client.project, dataset.dataset_id))
56+
return destination_data_set, destination_table_uri
57+
58+
59+
def test_ucaip_generated_batch_read_feature_values_sample_vision(capsys, shared_state):
60+
destination_data_set, destination_table_uri = setup_test()
61+
62+
featurestore_id = "perm_sample_featurestore"
63+
csv_read_instances = {"gcs_source": {"uris": [INPUT_CSV_FILE]}}
64+
destination = {"bigquery_destination": {"output_uri": destination_table_uri}}
65+
entity_type_specs = [
66+
{
67+
"entity_type_id": "perm_users",
68+
"feature_selector": {
69+
"id_matcher": {"ids": ["age", "gender", "liked_genres"]}
70+
},
71+
},
72+
{
73+
"entity_type_id": "perm_movies",
74+
"feature_selector": {"id_matcher": {"ids": ["average_rating", "genres"]}},
75+
},
76+
]
77+
78+
batch_read_feature_values_sample.batch_read_feature_values_sample(
79+
project=PROJECT_ID,
80+
featurestore_id=featurestore_id,
81+
csv_read_instances=csv_read_instances,
82+
destination=destination,
83+
entity_type_specs=entity_type_specs,
84+
)
85+
out, _ = capsys.readouterr()
86+
assert "batch_read_feature_values_response" in out
87+
with capsys.disabled():
88+
print(out)
89+
90+
shared_state["destination_data_set"] = destination_data_set

‎samples/snippets/feature_store_service/import_feature_values_sample.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START aiplatform_import_feature_values_sample]
16+
from google.cloud import aiplatform_v1beta1 as aiplatform
17+
18+
19+
def import_feature_values_sample(
20+
project: str,
21+
featurestore_id: str,
22+
entity_type_id: str,
23+
avro_source: str,
24+
feature_specs: str,
25+
entity_id_field: str,
26+
feature_time_field: str,
27+
worker_count: int = 2,
28+
location: str = "us-central1",
29+
api_endpoint: str = "us-central1-aiplatform.googleapis.com",
30+
timeout: int = 300,
31+
):
32+
# The AI Platform services require regional API endpoints.
33+
client_options = {"api_endpoint": api_endpoint}
34+
# Initialize client that will be used to create and send requests.
35+
# This client only needs to be created once, and can be reused for multiple requests.
36+
client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
37+
entity_type = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}"
38+
import_feature_values_request = {
39+
"entity_type": entity_type,
40+
"avro_source": avro_source,
41+
"feature_specs": feature_specs,
42+
"entity_id_field": entity_id_field,
43+
"feature_time_field": feature_time_field,
44+
"worker_count": worker_count,
45+
}
46+
lro_response = client.import_feature_values(request=import_feature_values_request)
47+
print("Long running operation:", lro_response.operation.name)
48+
import_feature_values_response = lro_response.result(timeout=timeout)
49+
print("import_feature_values_response:", import_feature_values_response)
50+
51+
52+
# [END aiplatform_import_feature_values_sample]

0 commit comments

Comments
 (0)