From bcf2fc1e72e1a125b4aac2f9f99fbfcf5e479993 Mon Sep 17 00:00:00 2001 From: Morgan Du Date: Thu, 27 Jan 2022 16:34:55 -0800 Subject: [PATCH 1/4] feat: add batch_serve_to_df, add unit tests, add integration tests; update: setup.py; fix: read_instances bug for batch_serve_to_* --- .../aiplatform/featurestore/featurestore.py | 380 +++++++++++++----- setup.py | 6 +- tests/system/aiplatform/test_featurestore.py | 86 +++- tests/unit/aiplatform/test_featurestores.py | 173 +++++--- 4 files changed, 480 insertions(+), 165 deletions(-) diff --git a/google/cloud/aiplatform/featurestore/featurestore.py b/google/cloud/aiplatform/featurestore/featurestore.py index 6d02bb9f76..e428aedc20 100644 --- a/google/cloud/aiplatform/featurestore/featurestore.py +++ b/google/cloud/aiplatform/featurestore/featurestore.py @@ -16,6 +16,7 @@ # from typing import Dict, List, Optional, Sequence, Tuple, Union +import uuid from google.auth import credentials as auth_credentials from google.protobuf import field_mask_pb2 @@ -32,6 +33,8 @@ from google.cloud.aiplatform import utils from google.cloud.aiplatform.utils import featurestore_utils +from google.cloud import bigquery + _LOGGER = base.Logger(__name__) @@ -591,6 +594,38 @@ def _batch_read_feature_values( return self + def _validate_and_get_read_instances( + self, read_instances_uri: str, + ) -> Union[gca_io.BigQuerySource, gca_io.CsvSource]: + """Gets read_instances + + Args: + read_instances_uri (str): + Required. Read_instances_uri can be either BigQuery URI to an input table, + or Google Cloud Storage URI to a csv file. + + Returns: + Union[gca_io.BigQuerySource, gca_io.CsvSource]: + BigQuery source or Csv source for read instances. The Csv source contains exactly 1 URI. + + Raises: + ValueError if read_instances_uri does not start with 'bq://' or 'gs://'. + """ + if not ( + read_instances_uri.startswith("bq://") + or read_instances_uri.startswith("gs://") + ): + raise ValueError( + "The read_instances_uri should be a single uri starts with either 'bq://' or 'gs://'." + ) + + if read_instances_uri.startswith("bq://"): + return gca_io.BigQuerySource(input_uri=read_instances_uri) + if read_instances_uri.startswith("gs://"): + return gca_io.CsvSource( + gcs_source=gca_io.GcsSource(uris=[read_instances_uri]) + ) + def _validate_and_get_batch_read_feature_values_request( self, serving_feature_ids: Dict[str, List[str]], @@ -599,9 +634,9 @@ def _validate_and_get_batch_read_feature_values_request( gca_io.CsvDestination, gca_io.TFRecordDestination, ], - feature_destination_fields: Optional[Dict[str, str]] = None, - read_instances: Optional[Union[gca_io.BigQuerySource, gca_io.CsvSource]] = None, + read_instances: Union[gca_io.BigQuerySource, gca_io.CsvSource], pass_through_fields: Optional[List[str]] = None, + feature_destination_fields: Optional[Dict[str, str]] = None, ) -> gca_featurestore_service.BatchReadFeatureValuesRequest: """Validates and gets batch_read_feature_values_request @@ -619,7 +654,17 @@ def _validate_and_get_batch_read_feature_values_request( destination (Union[gca_io.BigQueryDestination, gca_io.CsvDestination, gca_io.TFRecordDestination]): Required. BigQuery destination, Csv destination or TFRecord destination. + read_instances (Union[gca_io.BigQuerySource, gca_io.CsvSource]): + Required. BigQuery source or Csv source for read instances. + The Csv source must contain exactly 1 URI. + pass_through_fields (List[str]): + Optional. When not empty, the specified fields in the + read_instances source will be joined as-is in the output, + in addition to those fields from the Featurestore Entity. + For BigQuery source, the type of the pass-through values + will be automatically inferred. For CSV source, the + pass-through values will be passed as opaque bytes. feature_destination_fields (Dict[str, str]): Optional. A user defined dictionary to map a feature's fully qualified resource name to its destination field name. If the destination field name is not defined, @@ -631,17 +676,6 @@ def _validate_and_get_batch_read_feature_values_request( 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', } - read_instances (Union[gca_io.BigQuerySource, gca_io.CsvSource]): - Optional. BigQuery source or Csv source for read instances. - pass_through_fields (List[str]): - Optional. When not empty, the specified fields in the - read_instances source will be joined as-is in the output, - in addition to those fields from the Featurestore Entity. - - For BigQuery source, the type of the pass-through values - will be automatically inferred. For CSV source, the - pass-through values will be passed as opaque bytes. - Returns: gca_featurestore_service.BatchReadFeatureValuesRequest: batch read feature values request """ @@ -715,59 +749,14 @@ def _validate_and_get_batch_read_feature_values_request( return batch_read_feature_values_request - def _get_read_instances( - self, read_instances: Union[str, List[str]], - ) -> Union[gca_io.BigQuerySource, gca_io.CsvSource]: - """Gets read_instances - - Args: - read_instances (Union[str, List[str]]): - Required. Read_instances can be either BigQuery URI to the input table, - or Google Cloud Storage URI(-s) to the csv file(s). - - Returns: - Union[gca_io.BigQuerySource, gca_io.CsvSource]: BigQuery source or Csv source for read instances. - - Raises: - TypeError if read_instances is not a string or a list of strings. - ValueError if read_instances uri does not start with 'bq://' or 'gs://'. - ValueError if uris in read_instances do not start with 'gs://'. - """ - if isinstance(read_instances, str): - if not ( - read_instances.startswith("bq://") or read_instances.startswith("gs://") - ): - raise ValueError( - "The read_instances accepts a single uri starts with 'bq://' or 'gs://'." - ) - elif isinstance(read_instances, list) and all( - [isinstance(e, str) for e in read_instances] - ): - if not all([e.startswith("gs://") for e in read_instances]): - raise ValueError( - "The read_instances accepts a list of uris start with 'gs://' only." - ) - else: - raise TypeError( - "The read_instances type should to be either a str or a List[str]." - ) - - if isinstance(read_instances, str): - if read_instances.startswith("bq://"): - return gca_io.BigQuerySource(input_uri=read_instances) - else: - read_instances = [read_instances] - - return gca_io.CsvSource(gcs_source=gca_io.GcsSource(uris=read_instances)) - @base.optional_sync(return_input_arg="self") def batch_serve_to_bq( self, bq_destination_output_uri: str, serving_feature_ids: Dict[str, List[str]], - feature_destination_fields: Optional[Dict[str, str]] = None, - read_instances: Optional[Union[str, List[str]]] = None, + read_instances_uri: str, pass_through_fields: Optional[List[str]] = None, + feature_destination_fields: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, ) -> "Featurestore": @@ -793,29 +782,16 @@ def batch_serve_to_bq( 'my_entity_type_id_2': ['feature_id_2_1', 'feature_id_2_2'], } - feature_destination_fields (Dict[str, str]): - Optional. A user defined dictionary to map a feature's fully qualified resource name to - its destination field name. If the destination field name is not defined, - the feature ID will be used as its destination field name. + read_instances_uri (str): + Required. Read_instances_uri can be either BigQuery URI to an input table, + or Google Cloud Storage URI to a csv file. - Example: - feature_destination_fields = { - 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', - 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', - } - - read_instances (Union[str, List[str]]): - Optional. Read_instances can be either BigQuery URI to the input table, - or Google Cloud Storage URI(-s) to the - csv file(s). May contain wildcards. For more - information on wildcards, see - https://cloud.google.com/storage/docs/gsutil/addlhelp/WildcardNames. Example: 'bq://project.dataset.table_name' or - ["gs://my_bucket/my_file_1.csv", "gs://my_bucket/my_file_2.csv"] + "gs://my_bucket/my_file.csv" - Each read instance consists of exactly one read timestamp + Each read instance should consist of exactly one read timestamp and one or more entity IDs identifying entities of the corresponding EntityTypes whose Features are requested. @@ -844,6 +820,17 @@ def batch_serve_to_bq( will be automatically inferred. For CSV source, the pass-through values will be passed as opaque bytes. + feature_destination_fields (Dict[str, str]): + Optional. A user defined dictionary to map a feature's fully qualified resource name to + its destination field name. If the destination field name is not defined, + the feature ID will be used as its destination field name. + + Example: + feature_destination_fields = { + 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', + 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', + } + Returns: Featurestore: The featurestore resource object batch read feature values from. @@ -851,15 +838,15 @@ def batch_serve_to_bq( NotFound: if the BigQuery destination Dataset does not exist. FailedPrecondition: if the BigQuery destination Dataset/Table is in a different project. """ + read_instances = self._validate_and_get_read_instances(read_instances_uri) + batch_read_feature_values_request = self._validate_and_get_batch_read_feature_values_request( serving_feature_ids=serving_feature_ids, destination=gca_io.BigQueryDestination( output_uri=bq_destination_output_uri ), feature_destination_fields=feature_destination_fields, - read_instances=read_instances - if read_instances is None - else self._get_read_instances(read_instances), + read_instances=read_instances, pass_through_fields=pass_through_fields, ) @@ -874,9 +861,9 @@ def batch_serve_to_gcs( gcs_destination_output_uri_prefix: str, gcs_destination_type: str, serving_feature_ids: Dict[str, List[str]], - feature_destination_fields: Optional[Dict[str, str]] = None, - read_instances: Optional[Union[str, List[str]]] = None, + read_instances_uri: str, pass_through_fields: Optional[List[str]] = None, + feature_destination_fields: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, ) -> "Featurestore": @@ -923,29 +910,16 @@ def batch_serve_to_gcs( 'my_entity_type_id_2': ['feature_id_2_1', 'feature_id_2_2'], } - feature_destination_fields (Dict[str, str]): - Optional. A user defined dictionary to map a feature's fully qualified resource name to - its destination field name. If the destination field name is not defined, - the feature ID will be used as its destination field name. + read_instances_uri (str): + Required. Read_instances_uri can be either BigQuery URI to an input table, + or Google Cloud Storage URI to a csv file. - Example: - feature_destination_fields = { - 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', - 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', - } - - read_instances (Union[str, List[str]]): - Optional. Read_instances can be either BigQuery URI to the input table, - or Google Cloud Storage URI(-s) to the - csv file(s). May contain wildcards. For more - information on wildcards, see - https://cloud.google.com/storage/docs/gsutil/addlhelp/WildcardNames. Example: 'bq://project.dataset.table_name' or - ["gs://my_bucket/my_file_1.csv", "gs://my_bucket/my_file_2.csv"] + "gs://my_bucket/my_file.csv" - Each read instance consists of exactly one read timestamp + Each read instance should consist of exactly one read timestamp and one or more entity IDs identifying entities of the corresponding EntityTypes whose Features are requested. @@ -974,6 +948,17 @@ def batch_serve_to_gcs( will be automatically inferred. For CSV source, the pass-through values will be passed as opaque bytes. + feature_destination_fields (Dict[str, str]): + Optional. A user defined dictionary to map a feature's fully qualified resource name to + its destination field name. If the destination field name is not defined, + the feature ID will be used as its destination field name. + + Example: + feature_destination_fields = { + 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', + 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', + } + Returns: Featurestore: The featurestore resource object batch read feature values from. @@ -999,13 +984,13 @@ def batch_serve_to_gcs( if gcs_destination_type == "tfrecord": destination = gca_io.TFRecordDestination(gcs_destination=gcs_destination) + read_instances = self._validate_and_get_read_instances(read_instances_uri) + batch_read_feature_values_request = self._validate_and_get_batch_read_feature_values_request( serving_feature_ids=serving_feature_ids, destination=destination, feature_destination_fields=feature_destination_fields, - read_instances=read_instances - if read_instances is None - else self._get_read_instances(read_instances), + read_instances=read_instances, pass_through_fields=pass_through_fields, ) @@ -1013,3 +998,190 @@ def batch_serve_to_gcs( batch_read_feature_values_request=batch_read_feature_values_request, request_metadata=request_metadata, ) + + def batch_serve_to_df( + self, + serving_feature_ids: Dict[str, List[str]], + read_instances_df: "pd.DataFrame", # noqa: F821 - skip check for undefined name 'pd' + pass_through_fields: Optional[List[str]] = None, + feature_destination_fields: Optional[Dict[str, str]] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ) -> "pd.DataFrame": # noqa: F821 - skip check for undefined name 'pd' + """ Batch serves feature values to pandas DataFrame + + Note: + Calling this method will automatically create and delete a temporary + bigquery dataset in the same GCP project, which will be used + as the intermediary storage for batch serve feature values + from featurestore to dataframe. + + Args: + serving_feature_ids (Dict[str, List[str]]): + Required. A user defined dictionary to define the entity_types and their features for batch serve/read. + The keys of the dictionary are the serving entity_type ids and + the values are lists of serving feature ids in each entity_type. + + Example: + serving_feature_ids = { + 'my_entity_type_id_1': ['feature_id_1_1', 'feature_id_1_2'], + 'my_entity_type_id_2': ['feature_id_2_1', 'feature_id_2_2'], + } + + read_instances_df (pd.DataFrame): + Required. Read_instances_df is a pandas DataFrame containing the read instances. + + Each read instance should consist of exactly one read timestamp + and one or more entity IDs identifying entities of the + corresponding EntityTypes whose Features are requested. + + Each output instance contains Feature values of requested + entities concatenated together as of the read time. + + An example read_instances_df may be + pd.DataFrame( + data=[ + { + "my_entity_type_id_1": "my_entity_type_id_1_entity_1", + "my_entity_type_id_2": "my_entity_type_id_2_entity_1", + "timestamp": "2020-01-01T10:00:00.123Z" + ], + ) + + An example batch_serve_output_df may be + pd.DataFrame( + data=[ + { + "my_entity_type_id_1": "my_entity_type_id_1_entity_1", + "my_entity_type_id_2": "my_entity_type_id_2_entity_1", + "foo": "feature_id_1_1_feature_value", + "feature_id_1_2": "feature_id_1_2_feature_value", + "feature_id_2_1": "feature_id_2_1_feature_value", + "bar": "feature_id_2_2_feature_value", + "timestamp": "2020-01-01T10:00:00.123Z" + ], + ) + + Timestamp in each read instance must be millisecond-aligned. + + The columns can be in any order. + + Values in the timestamp column must use the RFC 3339 format, + e.g. ``2012-07-30T10:43:17.123Z``. + + pass_through_fields (List[str]): + Optional. When not empty, the specified fields in the + read_instances source will be joined as-is in the output, + in addition to those fields from the Featurestore Entity. + + For BigQuery source, the type of the pass-through values + will be automatically inferred. For CSV source, the + pass-through values will be passed as opaque bytes. + + feature_destination_fields (Dict[str, str]): + Optional. A user defined dictionary to map a feature's fully qualified resource name to + its destination field name. If the destination field name is not defined, + the feature ID will be used as its destination field name. + + Example: + feature_destination_fields = { + 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', + 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', + } + + Returns: + pd.DataFrame: The pandas DataFrame containing feature values from batch serving. + + """ + try: + from google.cloud import bigquery_storage + except ImportError: + raise ImportError( + f"Google-Cloud-Bigquery-Storage is not installed. Please install google-cloud-bigquery-storage to use " + f"{self.batch_serve_to_df.__name__}" + ) + + try: + import pyarrow # noqa: F401 - skip check for 'pyarrow' which is required when using 'google.cloud.bigquery' + except ImportError: + raise ImportError( + f"Pyarrow is not installed. Please install pyarrow to use " + f"{self.batch_serve_to_df.__name__}" + ) + + try: + import pandas as pd + except ImportError: + raise ImportError( + f"Pandas is not installed. Please install pandas to use " + f"{self.batch_serve_to_df.__name__}" + ) + + bigquery_client = bigquery.Client( + project=self.project, credentials=self.credentials + ) + + featurestore_name_components = self._parse_resource_name(self.resource_name) + featurestore_id = featurestore_name_components["featurestore"] + + temp_bq_dataset_name = f"temp_{featurestore_id}_{uuid.uuid4()}".replace( + "-", "_" + ) + temp_bq_dataset_id = f"{initializer.global_config.project}.{temp_bq_dataset_name}"[ + :1024 + ] + temp_bq_dataset = bigquery.Dataset(dataset_ref=temp_bq_dataset_id) + temp_bq_dataset.location = self.location + temp_bq_dataset = bigquery_client.create_dataset(temp_bq_dataset) + + temp_bq_batch_serve_table_name = "batch_serve" + temp_bq_read_instances_table_name = "read_instances" + temp_bq_batch_serve_table_id = ( + f"{temp_bq_dataset_id}.{temp_bq_batch_serve_table_name}" + ) + temp_bq_read_instances_table_id = ( + f"{temp_bq_dataset_id}.{temp_bq_read_instances_table_name}" + ) + + try: + + job = bigquery_client.load_table_from_dataframe( + dataframe=read_instances_df, destination=temp_bq_read_instances_table_id + ) + job.result() + + self.batch_serve_to_bq( + bq_destination_output_uri=f"bq://{temp_bq_batch_serve_table_id}", + serving_feature_ids=serving_feature_ids, + read_instances_uri=f"bq://{temp_bq_read_instances_table_id}", + pass_through_fields=pass_through_fields, + feature_destination_fields=feature_destination_fields, + request_metadata=request_metadata, + ) + + bigquery_storage_read_client = bigquery_storage.BigQueryReadClient( + credentials=self.credentials + ) + read_session_proto = bigquery_storage_read_client.create_read_session( + parent=f"projects/{self.project}", + read_session=bigquery_storage.types.ReadSession( + table="projects/{project}/datasets/{dataset}/tables/{table}".format( + project=self.project, + dataset=temp_bq_dataset_name, + table=temp_bq_batch_serve_table_name, + ), + data_format=bigquery_storage.types.DataFormat.ARROW, + ), + ) + + finally: + bigquery_client.delete_dataset( + dataset=temp_bq_dataset.dataset_id, delete_contents=True, + ) + + frames = [] + for stream in read_session_proto.streams: + reader = bigquery_storage_read_client.read_rows(stream.name) + for message in reader.rows().pages: + frames.append(message.to_dataframe()) + + return None if not frames else pd.concat(frames) diff --git a/setup.py b/setup.py index 011da8aea8..8f8b297c1b 100644 --- a/setup.py +++ b/setup.py @@ -47,7 +47,11 @@ "werkzeug >= 2.0.0", "tensorflow >=2.4.0", ] -featurestore_extra_require = ["pandas >= 1.0.0", "pyarrow >= 6.0.1"] +featurestore_extra_require = [ + "google-cloud-bigquery-storage >= 2.11.0", + "pandas >= 1.0.0", + "pyarrow >= 6.0.1", +] full_extra_require = list( set( diff --git a/tests/system/aiplatform/test_featurestore.py b/tests/system/aiplatform/test_featurestore.py index 03070eee3c..cbbfd82efb 100644 --- a/tests/system/aiplatform/test_featurestore.py +++ b/tests/system/aiplatform/test_featurestore.py @@ -411,6 +411,82 @@ def test_search_features(self, shared_state): len(list_searched_features) - shared_state["base_list_searched_features"] ) == 6 + def test_batch_serve_to_df(self, shared_state, caplog): + + assert shared_state["featurestore"] + assert shared_state["user_age_feature_resource_name"] + assert shared_state["user_gender_feature_resource_name"] + assert shared_state["user_liked_genres_feature_resource_name"] + + featurestore = shared_state["featurestore"] + + user_age_feature_resource_name = shared_state["user_age_feature_resource_name"] + user_gender_feature_resource_name = shared_state[ + "user_gender_feature_resource_name" + ] + user_liked_genres_feature_resource_name = shared_state[ + "user_liked_genres_feature_resource_name" + ] + + aiplatform.init( + project=e2e_base._PROJECT, location=e2e_base._LOCATION, + ) + + caplog.set_level(logging.INFO) + + read_instances_df = pd.DataFrame( + data=[ + ["alice", "movie_01", "2021-09-15T08:28:14Z"], + ["bob", "movie_02", "2021-09-15T08:28:14Z"], + ["dav", "movie_03", "2021-09-15T08:28:14Z"], + ["eve", "movie_04", "2021-09-15T08:28:14Z"], + ["alice", "movie_03", "2021-09-14T09:35:15Z"], + ["bob", "movie_04", "2020-02-14T09:35:15Z"], + ], + columns=["users", "movies", "timestamp"], + ) + read_instances_df = read_instances_df.astype({"timestamp": "datetime64"}) + + df = featurestore.batch_serve_to_df( + serving_feature_ids={ + _TEST_USER_ENTITY_TYPE_ID: [ + _TEST_USER_AGE_FEATURE_ID, + _TEST_USER_GENDER_FEATURE_ID, + _TEST_USER_LIKED_GENRES_FEATURE_ID, + ], + _TEST_MOVIE_ENTITY_TYPE_ID: [ + _TEST_MOVIE_TITLE_FEATURE_ID, + _TEST_MOVIE_GENRES_FEATURE_ID, + _TEST_MOVIE_AVERAGE_RATING_FEATURE_ID, + ], + }, + read_instances_df=read_instances_df, + feature_destination_fields={ + user_age_feature_resource_name: "user_age_dest", + user_gender_feature_resource_name: "user_gender_dest", + user_liked_genres_feature_resource_name: "user_liked_genres_dest", + }, + ) + + expected_df_columns = [ + "timestamp", + "entity_type_users", + "user_age_dest", + "user_gender_dest", + "user_liked_genres_dest", + "entity_type_movies", + "title", + "genres", + "average_rating", + ] + + assert type(df) == pd.DataFrame + assert list(df.columns) == expected_df_columns + assert df.size == 54 + assert "Featurestore feature values served." in caplog.text + + caplog.clear() + def test_batch_serve_to_gcs(self, shared_state, caplog): assert shared_state["featurestore"] @@ -448,12 +524,12 @@ def test_batch_serve_to_gcs(self, shared_state, caplog): _TEST_MOVIE_AVERAGE_RATING_FEATURE_ID, ], }, + read_instances_uri=_TEST_READ_INSTANCE_SRC, feature_destination_fields={ user_age_feature_resource_name: "user_age_dest", user_gender_feature_resource_name: "user_gender_dest", user_liked_genres_feature_resource_name: "user_liked_genres_dest", }, - read_instances=_TEST_READ_INSTANCE_SRC, gcs_destination_output_uri_prefix=f"gs://{bucket_name}/featurestore_test/tfrecord", gcs_destination_type="tfrecord", ) @@ -498,12 +574,12 @@ def test_batch_serve_to_bq(self, shared_state, caplog): _TEST_MOVIE_AVERAGE_RATING_FEATURE_ID, ], }, + read_instances_uri=_TEST_READ_INSTANCE_SRC, feature_destination_fields={ user_age_feature_resource_name: "user_age_dest", user_gender_feature_resource_name: "user_gender_dest", user_liked_genres_feature_resource_name: "user_liked_genres_dest", }, - read_instances=_TEST_READ_INSTANCE_SRC, bq_destination_output_uri=f"bq://{bigquery_dataset_id}.test_table", ) @@ -525,3 +601,9 @@ def test_online_reads(self, shared_state): feature_ids=[_TEST_MOVIE_TITLE_FEATURE_ID, _TEST_MOVIE_GENRES_FEATURE_ID], ) assert type(movie_entity_views) == pd.DataFrame + + movie_entity_views = movie_entity_type.read( + entity_ids="movie_01", + feature_ids=[_TEST_MOVIE_TITLE_FEATURE_ID, _TEST_MOVIE_GENRES_FEATURE_ID], + ) + assert type(movie_entity_views) == pd.DataFrame diff --git a/tests/unit/aiplatform/test_featurestores.py b/tests/unit/aiplatform/test_featurestores.py index 449d0348c3..5da494c443 100644 --- a/tests/unit/aiplatform/test_featurestores.py +++ b/tests/unit/aiplatform/test_featurestores.py @@ -53,6 +53,9 @@ ) from google.cloud import bigquery +from google.cloud import bigquery_storage + +from google.cloud.bigquery_storage_v1.types import stream as gcbqs_stream # project _TEST_PROJECT = "test-project" @@ -379,6 +382,30 @@ def bq_delete_dataset_mock(bq_init_client_mock): yield bq_delete_dataset_mock +@pytest.fixture +def bqs_client_mock(): + mock = MagicMock(bigquery_storage.BigQueryReadClient) + yield mock + + +@pytest.fixture +def bqs_init_client_mock(bqs_client_mock): + with patch.object(bigquery_storage, "BigQueryReadClient") as bqs_init_client_mock: + bqs_init_client_mock.return_value = bqs_client_mock + yield bqs_init_client_mock + + +@pytest.fixture +def bqs_create_read_session(bqs_init_client_mock): + with patch.object( + bigquery_storage.BigQueryReadClient, "create_read_session" + ) as bqs_create_read_session: + read_session_proto = gcbqs_stream.ReadSession() + read_session_proto.streams = [gcbqs_stream.ReadStream()] + bqs_create_read_session.return_value = read_session_proto + yield bqs_create_read_session + + # All Featurestore Mocks @pytest.fixture def get_featurestore_mock(): @@ -1069,40 +1096,6 @@ def test_validate_and_get_batch_read_feature_values_request( my_featurestore = aiplatform.Featurestore( featurestore_name=_TEST_FEATURESTORE_NAME ) - expected_batch_read_feature_values_request = gca_featurestore_service.BatchReadFeatureValuesRequest( - featurestore=my_featurestore.resource_name, - destination=gca_featurestore_service.FeatureValueDestination( - bigquery_destination=_TEST_BQ_DESTINATION, - ), - entity_type_specs=expected_entity_type_specs, - ) - assert ( - expected_batch_read_feature_values_request - == my_featurestore._validate_and_get_batch_read_feature_values_request( - serving_feature_ids=serving_feature_ids, - destination=_TEST_BQ_DESTINATION, - feature_destination_fields=feature_destination_fields, - ) - ) - - @pytest.mark.usefixtures("get_featurestore_mock") - def test_validate_and_get_batch_read_feature_values_request_with_read_instances( - self, - ): - aiplatform.init(project=_TEST_PROJECT) - my_featurestore = aiplatform.Featurestore( - featurestore_name=_TEST_FEATURESTORE_NAME - ) - expected_entity_type_specs = [ - _get_entity_type_spec_proto_with_feature_ids( - entity_type_id="my_entity_type_id_1", - feature_ids=["my_feature_id_1_1", "my_feature_id_1_2"], - ), - _get_entity_type_spec_proto_with_feature_ids( - entity_type_id="my_entity_type_id_2", - feature_ids=["my_feature_id_2_1", "my_feature_id_2_2"], - ), - ] expected_batch_read_feature_values_request = gca_featurestore_service.BatchReadFeatureValuesRequest( featurestore=my_featurestore.resource_name, destination=gca_featurestore_service.FeatureValueDestination( @@ -1114,60 +1107,53 @@ def test_validate_and_get_batch_read_feature_values_request_with_read_instances( assert ( expected_batch_read_feature_values_request == my_featurestore._validate_and_get_batch_read_feature_values_request( - serving_feature_ids=_TEST_SERVING_FEATURE_IDS, + serving_feature_ids=serving_feature_ids, destination=_TEST_BQ_DESTINATION, read_instances=_TEST_BQ_SOURCE, + feature_destination_fields=feature_destination_fields, ) ) @pytest.mark.usefixtures("get_featurestore_mock") @pytest.mark.parametrize( - "read_instances, expected", + "read_instances_uri, expected_read_instances", [ (_TEST_BQ_SOURCE_URI, _TEST_BQ_SOURCE), - (_TEST_GCS_CSV_SOURCE_URIS, _TEST_CSV_SOURCE), (_TEST_GCS_CSV_SOURCE_URI, _TEST_CSV_SOURCE), ], ) - def test_get_read_instances(self, read_instances, expected): + def test_validate_and_get_read_instances( + self, read_instances_uri, expected_read_instances + ): aiplatform.init(project=_TEST_PROJECT) my_featurestore = aiplatform.Featurestore( featurestore_name=_TEST_FEATURESTORE_NAME ) - assert expected == my_featurestore._get_read_instances( - read_instances=read_instances - ) - - @pytest.mark.usefixtures("get_featurestore_mock") - @pytest.mark.parametrize( - "read_instances", - [[1, 2, 3, 4, 5], 1, (_TEST_GCS_CSV_SOURCE_URI, _TEST_GCS_CSV_SOURCE_URI)], - ) - def test_get_read_instances_with_raise_typeerror(self, read_instances): - aiplatform.init(project=_TEST_PROJECT) - my_featurestore = aiplatform.Featurestore( - featurestore_name=_TEST_FEATURESTORE_NAME + assert ( + expected_read_instances + == my_featurestore._validate_and_get_read_instances( + read_instances_uri=read_instances_uri + ) ) - with pytest.raises(TypeError): - my_featurestore._get_read_instances(read_instances=read_instances) @pytest.mark.usefixtures("get_featurestore_mock") @pytest.mark.parametrize( - "read_instances", + "read_instances_uri", [ "gcs://my_bucket/my_file_1.csv", - "bigquery://my_bucket/my_file_1.csv", + "bigquery://project.dataset.table_name", "my_bucket/my_file_1.csv", - [_TEST_BQ_SOURCE_URI], ], ) - def test_get_read_instances_with_raise_valueerror(self, read_instances): + def test_validate_and_get_read_instances_with_raise(self, read_instances_uri): aiplatform.init(project=_TEST_PROJECT) my_featurestore = aiplatform.Featurestore( featurestore_name=_TEST_FEATURESTORE_NAME ) with pytest.raises(ValueError): - my_featurestore._get_read_instances(read_instances=read_instances) + my_featurestore._validate_and_get_read_instances( + read_instances_uri=read_instances_uri + ) @pytest.mark.parametrize("sync", [True, False]) @pytest.mark.usefixtures("get_featurestore_mock") @@ -1194,11 +1180,13 @@ def test_batch_serve_to_bq(self, batch_read_feature_values_mock, sync): bigquery_destination=_TEST_BQ_DESTINATION, ), entity_type_specs=expected_entity_type_specs, + bigquery_read_instances=_TEST_BQ_SOURCE, ) my_featurestore.batch_serve_to_bq( bq_destination_output_uri=_TEST_BQ_DESTINATION_URI, serving_feature_ids=_TEST_SERVING_FEATURE_IDS, + read_instances_uri=_TEST_BQ_SOURCE_URI, sync=sync, ) @@ -1235,12 +1223,14 @@ def test_batch_serve_to_gcs(self, batch_read_feature_values_mock, sync): tfrecord_destination=_TEST_TFRECORD_DESTINATION, ), entity_type_specs=expected_entity_type_specs, + csv_read_instances=_TEST_CSV_SOURCE, ) my_featurestore.batch_serve_to_gcs( gcs_destination_output_uri_prefix=_TEST_GCS_OUTPUT_URI_PREFIX, gcs_destination_type=_TEST_GCS_DESTINATION_TYPE_TFRECORD, serving_feature_ids=_TEST_SERVING_FEATURE_IDS, + read_instances_uri=_TEST_GCS_CSV_SOURCE_URI, sync=sync, ) @@ -1265,8 +1255,75 @@ def test_batch_serve_to_gcs_with_invalid_gcs_destination_type(self): gcs_destination_output_uri_prefix=_TEST_GCS_OUTPUT_URI_PREFIX, gcs_destination_type=_TEST_GCS_DESTINATION_TYPE_INVALID, serving_feature_ids=_TEST_SERVING_FEATURE_IDS, + read_instances_uri=_TEST_GCS_CSV_SOURCE_URI, ) + @pytest.mark.usefixtures( + "get_featurestore_mock", + "bq_init_client_mock", + "bq_init_dataset_mock", + "bq_create_dataset_mock", + "bq_load_table_from_dataframe_mock", + "bq_delete_dataset_mock", + "bqs_init_client_mock", + "bqs_create_read_session", + ) + @patch("uuid.uuid4", uuid_mock) + def test_batch_serve_to_df(self, batch_read_feature_values_mock): + aiplatform.init(project=_TEST_PROJECT) + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_NAME + ) + + read_instances_df = pd.DataFrame() + + expected_temp_bq_dataset_name = f"temp_{_TEST_FEATURESTORE_ID}_{uuid.uuid4()}".replace( + "-", "_" + ) + expecte_temp_bq_dataset_id = f"{initializer.global_config.project}.{expected_temp_bq_dataset_name}"[ + :1024 + ] + expected_temp_bq_read_instances_table_id = ( + f"{expecte_temp_bq_dataset_id}.read_instances" + ) + expected_temp_bq_batch_serve_table_id = ( + f"{expecte_temp_bq_dataset_id}.batch_serve" + ) + + expected_entity_type_specs = [ + _get_entity_type_spec_proto_with_feature_ids( + entity_type_id="my_entity_type_id_1", + feature_ids=["my_feature_id_1_1", "my_feature_id_1_2"], + ), + _get_entity_type_spec_proto_with_feature_ids( + entity_type_id="my_entity_type_id_2", + feature_ids=["my_feature_id_2_1", "my_feature_id_2_2"], + ), + ] + + expected_batch_read_feature_values_request = gca_featurestore_service.BatchReadFeatureValuesRequest( + featurestore=my_featurestore.resource_name, + destination=gca_featurestore_service.FeatureValueDestination( + bigquery_destination=gca_io.BigQueryDestination( + output_uri=f"bq://{expected_temp_bq_batch_serve_table_id}" + ), + ), + entity_type_specs=expected_entity_type_specs, + bigquery_read_instances=gca_io.BigQuerySource( + input_uri=f"bq://{expected_temp_bq_read_instances_table_id}" + ), + ) + + my_featurestore.batch_serve_to_df( + serving_feature_ids=_TEST_SERVING_FEATURE_IDS, + read_instances_df=read_instances_df, + ) + + batch_read_feature_values_mock.assert_called_once_with( + request=expected_batch_read_feature_values_request, + metadata=_TEST_REQUEST_METADATA, + ) + class TestEntityType: def setup_method(self): From 4b03f01e11fbefa44dd644796d3b89c9ffe0ba74 Mon Sep 17 00:00:00 2001 From: Morgan Du Date: Mon, 31 Jan 2022 17:46:01 -0800 Subject: [PATCH 2/4] fix: add self.wait() before self._parse_resource_name, add ignore_index=True for pd.concat and return pd.DataFrame upon empty frames --- .../aiplatform/featurestore/entity_type.py | 5 +++++ .../cloud/aiplatform/featurestore/feature.py | 2 ++ .../aiplatform/featurestore/featurestore.py | 19 ++++++++++++------- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/google/cloud/aiplatform/featurestore/entity_type.py b/google/cloud/aiplatform/featurestore/entity_type.py index 8a85b1aa7a..67b5b9d4ab 100644 --- a/google/cloud/aiplatform/featurestore/entity_type.py +++ b/google/cloud/aiplatform/featurestore/entity_type.py @@ -126,6 +126,7 @@ def __init__( @property def featurestore_name(self) -> str: """Full qualified resource name of the managed featurestore in which this EntityType is.""" + self.wait() entity_type_name_components = self._parse_resource_name(self.resource_name) return featurestore.Featurestore._format_resource_name( project=entity_type_name_components["project"], @@ -150,6 +151,7 @@ def get_feature(self, feature_id: str) -> "featurestore.Feature": Returns: featurestore.Feature - The managed feature resource object. """ + self.wait() entity_type_name_components = self._parse_resource_name(self.resource_name) return featurestore.Feature( @@ -1213,6 +1215,7 @@ def ingest_from_df( project=self.project, credentials=self.credentials ) + self.wait() entity_type_name_components = self._parse_resource_name(self.resource_name) featurestore_id, entity_type_id = ( entity_type_name_components["featurestore"], @@ -1222,6 +1225,8 @@ def ingest_from_df( temp_bq_dataset_name = f"temp_{featurestore_id}_{uuid.uuid4()}".replace( "-", "_" ) + + # TODO(b/216497263): Add support for resource project does not match initializer.global_config.project temp_bq_dataset_id = f"{initializer.global_config.project}.{temp_bq_dataset_name}"[ :1024 ] diff --git a/google/cloud/aiplatform/featurestore/feature.py b/google/cloud/aiplatform/featurestore/feature.py index 08c91cdb1a..3c623d6b9e 100644 --- a/google/cloud/aiplatform/featurestore/feature.py +++ b/google/cloud/aiplatform/featurestore/feature.py @@ -125,6 +125,7 @@ def __init__( @property def featurestore_name(self) -> str: """Full qualified resource name of the managed featurestore in which this Feature is.""" + self.wait() feature_path_components = self._parse_resource_name(self.resource_name) return featurestore.Featurestore._format_resource_name( @@ -144,6 +145,7 @@ def get_featurestore(self) -> "featurestore.Featurestore": @property def entity_type_name(self) -> str: """Full qualified resource name of the managed entityType in which this Feature is.""" + self.wait() feature_path_components = self._parse_resource_name(self.resource_name) return featurestore.EntityType._format_resource_name( diff --git a/google/cloud/aiplatform/featurestore/featurestore.py b/google/cloud/aiplatform/featurestore/featurestore.py index e428aedc20..a9504ee498 100644 --- a/google/cloud/aiplatform/featurestore/featurestore.py +++ b/google/cloud/aiplatform/featurestore/featurestore.py @@ -112,6 +112,7 @@ def get_entity_type(self, entity_type_id: str) -> "featurestore.EntityType": Returns: featurestore.EntityType - The managed entityType resource object. """ + self.wait() featurestore_name_components = self._parse_resource_name(self.resource_name) return featurestore.EntityType( @@ -680,6 +681,7 @@ def _validate_and_get_batch_read_feature_values_request( gca_featurestore_service.BatchReadFeatureValuesRequest: batch read feature values request """ + self.wait() featurestore_name_components = self._parse_resource_name(self.resource_name) feature_destination_fields = feature_destination_fields or {} @@ -1120,12 +1122,15 @@ def batch_serve_to_df( project=self.project, credentials=self.credentials ) + self.wait() featurestore_name_components = self._parse_resource_name(self.resource_name) featurestore_id = featurestore_name_components["featurestore"] temp_bq_dataset_name = f"temp_{featurestore_id}_{uuid.uuid4()}".replace( "-", "_" ) + + # TODO(b/216497263): Add support for resource project does not match initializer.global_config.project temp_bq_dataset_id = f"{initializer.global_config.project}.{temp_bq_dataset_name}"[ :1024 ] @@ -1173,15 +1178,15 @@ def batch_serve_to_df( ), ) + frames = [] + for stream in read_session_proto.streams: + reader = bigquery_storage_read_client.read_rows(stream.name) + for message in reader.rows().pages: + frames.append(message.to_dataframe()) + finally: bigquery_client.delete_dataset( dataset=temp_bq_dataset.dataset_id, delete_contents=True, ) - frames = [] - for stream in read_session_proto.streams: - reader = bigquery_storage_read_client.read_rows(stream.name) - for message in reader.rows().pages: - frames.append(message.to_dataframe()) - - return None if not frames else pd.concat(frames) + return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame(frames) From e7467e959b66f4f8173e2318be12213a2c293a34 Mon Sep 17 00:00:00 2001 From: Morgan Du Date: Tue, 1 Feb 2022 13:38:56 -0800 Subject: [PATCH 3/4] fix: setup.py versioning conflict --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 8f8b297c1b..8c1872c300 100644 --- a/setup.py +++ b/setup.py @@ -48,7 +48,7 @@ "tensorflow >=2.4.0", ] featurestore_extra_require = [ - "google-cloud-bigquery-storage >= 2.11.0", + "google-cloud-bigquery-storage", "pandas >= 1.0.0", "pyarrow >= 6.0.1", ] From 290b07d70c0a01300e79507746a06a9c0dcb8d10 Mon Sep 17 00:00:00 2001 From: Morgan Du Date: Thu, 3 Feb 2022 21:17:27 -0800 Subject: [PATCH 4/4] fix: optional_sync --- .../aiplatform/featurestore/entity_type.py | 46 ++++++++++++++----- .../cloud/aiplatform/featurestore/feature.py | 27 +++++++---- .../aiplatform/featurestore/featurestore.py | 31 ++++++++++--- tests/unit/aiplatform/test_featurestores.py | 5 +- 4 files changed, 79 insertions(+), 30 deletions(-) diff --git a/google/cloud/aiplatform/featurestore/entity_type.py b/google/cloud/aiplatform/featurestore/entity_type.py index 67b5b9d4ab..274f89d2aa 100644 --- a/google/cloud/aiplatform/featurestore/entity_type.py +++ b/google/cloud/aiplatform/featurestore/entity_type.py @@ -123,10 +123,8 @@ def __init__( location=self.location, credentials=credentials, ) - @property - def featurestore_name(self) -> str: - """Full qualified resource name of the managed featurestore in which this EntityType is.""" - self.wait() + def _get_featurestore_name(self) -> str: + """Gets full qualified resource name of the managed featurestore in which this EntityType is.""" entity_type_name_components = self._parse_resource_name(self.resource_name) return featurestore.Featurestore._format_resource_name( project=entity_type_name_components["project"], @@ -134,6 +132,12 @@ def featurestore_name(self) -> str: featurestore=entity_type_name_components["featurestore"], ) + @property + def featurestore_name(self) -> str: + """Full qualified resource name of the managed featurestore in which this EntityType is.""" + self.wait() + return self._get_featurestore_name() + def get_featurestore(self) -> "featurestore.Featurestore": """Retrieves the managed featurestore in which this EntityType is. @@ -142,7 +146,7 @@ def get_featurestore(self) -> "featurestore.Featurestore": """ return featurestore.Featurestore(self.featurestore_name) - def get_feature(self, feature_id: str) -> "featurestore.Feature": + def _get_feature(self, feature_id: str) -> "featurestore.Feature": """Retrieves an existing managed feature in this EntityType. Args: @@ -151,9 +155,7 @@ def get_feature(self, feature_id: str) -> "featurestore.Feature": Returns: featurestore.Feature - The managed feature resource object. """ - self.wait() entity_type_name_components = self._parse_resource_name(self.resource_name) - return featurestore.Feature( feature_name=featurestore.Feature._format_resource_name( project=entity_type_name_components["project"], @@ -164,6 +166,18 @@ def get_feature(self, feature_id: str) -> "featurestore.Feature": ) ) + def get_feature(self, feature_id: str) -> "featurestore.Feature": + """Retrieves an existing managed feature in this EntityType. + + Args: + feature_id (str): + Required. The managed feature resource ID in this EntityType. + Returns: + featurestore.Feature - The managed feature resource object. + """ + self.wait() + return self._get_feature(feature_id=feature_id) + def update( self, description: Optional[str] = None, @@ -204,6 +218,7 @@ def update( Returns: EntityType - The updated entityType resource object. """ + self.wait() update_mask = list() if description: @@ -382,6 +397,7 @@ def list_features( Returns: List[featurestore.Feature] - A list of managed feature resource objects. """ + self.wait() return featurestore.Feature.list( entity_type_name=self.resource_name, filter=filter, order_by=order_by, ) @@ -401,7 +417,7 @@ def delete_features(self, feature_ids: List[str], sync: bool = True,) -> None: """ features = [] for feature_id in feature_ids: - feature = self.get_feature(feature_id=feature_id) + feature = self._get_feature(feature_id=feature_id) feature.delete(sync=False) features.append(feature) @@ -628,6 +644,7 @@ def create_feature( featurestore.Feature - feature resource object """ + self.wait() return featurestore.Feature.create( feature_id=feature_id, value_type=value_type, @@ -763,8 +780,9 @@ def batch_create_features( return self + @staticmethod def _validate_and_get_import_feature_values_request( - self, + entity_type_name: str, feature_ids: List[str], feature_time: Union[str, datetime.datetime], data_source: Union[gca_io.AvroSource, gca_io.BigQuerySource, gca_io.CsvSource], @@ -775,6 +793,8 @@ def _validate_and_get_import_feature_values_request( ) -> gca_featurestore_service.ImportFeatureValuesRequest: """Validates and get import feature values request. Args: + entity_type_name (str): + Required. A fully-qualified entityType resource name. feature_ids (List[str]): Required. IDs of the Feature to import values of. The Features must exist in the target @@ -842,7 +862,7 @@ def _validate_and_get_import_feature_values_request( ] import_feature_values_request = gca_featurestore_service.ImportFeatureValuesRequest( - entity_type=self.resource_name, + entity_type=entity_type_name, feature_specs=feature_specs, entity_id_field=entity_id_field, disable_online_serving=disable_online_serving, @@ -994,6 +1014,7 @@ def ingest_from_bq( bigquery_source = gca_io.BigQuerySource(input_uri=bq_source_uri) import_feature_values_request = self._validate_and_get_import_feature_values_request( + entity_type_name=self.resource_name, feature_ids=feature_ids, feature_time=feature_time, data_source=bigquery_source, @@ -1116,6 +1137,7 @@ def ingest_from_gcs( data_source = gca_io.AvroSource(gcs_source=gcs_source) import_feature_values_request = self._validate_and_get_import_feature_values_request( + entity_type_name=self.resource_name, feature_ids=feature_ids, feature_time=feature_time, data_source=data_source, @@ -1302,7 +1324,7 @@ def read( Returns: pd.DataFrame: entities' feature values in DataFrame """ - + self.wait() if isinstance(feature_ids, str): feature_ids = [feature_ids] @@ -1344,7 +1366,7 @@ def read( feature_descriptor.id for feature_descriptor in header.feature_descriptors ] - return EntityType._construct_dataframe( + return self._construct_dataframe( feature_ids=feature_ids, entity_views=entity_views, ) diff --git a/google/cloud/aiplatform/featurestore/feature.py b/google/cloud/aiplatform/featurestore/feature.py index 3c623d6b9e..1564abfd62 100644 --- a/google/cloud/aiplatform/featurestore/feature.py +++ b/google/cloud/aiplatform/featurestore/feature.py @@ -122,18 +122,21 @@ def __init__( else featurestore_id, ) - @property - def featurestore_name(self) -> str: - """Full qualified resource name of the managed featurestore in which this Feature is.""" - self.wait() + def _get_featurestore_name(self) -> str: + """Gets full qualified resource name of the managed featurestore in which this Feature is.""" feature_path_components = self._parse_resource_name(self.resource_name) - return featurestore.Featurestore._format_resource_name( project=feature_path_components["project"], location=feature_path_components["location"], featurestore=feature_path_components["featurestore"], ) + @property + def featurestore_name(self) -> str: + """Full qualified resource name of the managed featurestore in which this Feature is.""" + self.wait() + return self._get_featurestore_name() + def get_featurestore(self) -> "featurestore.Featurestore": """Retrieves the managed featurestore in which this Feature is. @@ -142,12 +145,9 @@ def get_featurestore(self) -> "featurestore.Featurestore": """ return featurestore.Featurestore(featurestore_name=self.featurestore_name) - @property - def entity_type_name(self) -> str: - """Full qualified resource name of the managed entityType in which this Feature is.""" - self.wait() + def _get_entity_type_name(self) -> str: + """Gets full qualified resource name of the managed entityType in which this Feature is.""" feature_path_components = self._parse_resource_name(self.resource_name) - return featurestore.EntityType._format_resource_name( project=feature_path_components["project"], location=feature_path_components["location"], @@ -155,6 +155,12 @@ def entity_type_name(self) -> str: entity_type=feature_path_components["entity_type"], ) + @property + def entity_type_name(self) -> str: + """Full qualified resource name of the managed entityType in which this Feature is.""" + self.wait() + return self._get_entity_type_name() + def get_entity_type(self) -> "featurestore.EntityType": """Retrieves the managed entityType in which this Feature is. @@ -205,6 +211,7 @@ def update( Returns: Feature - The updated feature resource object. """ + self.wait() update_mask = list() if description: diff --git a/google/cloud/aiplatform/featurestore/featurestore.py b/google/cloud/aiplatform/featurestore/featurestore.py index a9504ee498..648dd7ad2b 100644 --- a/google/cloud/aiplatform/featurestore/featurestore.py +++ b/google/cloud/aiplatform/featurestore/featurestore.py @@ -113,8 +113,18 @@ def get_entity_type(self, entity_type_id: str) -> "featurestore.EntityType": featurestore.EntityType - The managed entityType resource object. """ self.wait() - featurestore_name_components = self._parse_resource_name(self.resource_name) + return self._get_entity_type(entity_type_id=entity_type_id) + + def _get_entity_type(self, entity_type_id: str) -> "featurestore.EntityType": + """Retrieves an existing managed entityType in this Featurestore. + Args: + entity_type_id (str): + Required. The managed entityType resource ID in this Featurestore. + Returns: + featurestore.EntityType - The managed entityType resource object. + """ + featurestore_name_components = self._parse_resource_name(self.resource_name) return featurestore.EntityType( entity_type_name=featurestore.EntityType._format_resource_name( project=featurestore_name_components["project"], @@ -225,6 +235,7 @@ def _update( Returns: Featurestore - The updated featurestore resource object. """ + self.wait() update_mask = list() if labels: @@ -314,6 +325,7 @@ def list_entity_types( Returns: List[featurestore.EntityType] - A list of managed entityType resource objects. """ + self.wait() return featurestore.EntityType.list( featurestore_name=self.resource_name, filter=filter, order_by=order_by, ) @@ -338,7 +350,7 @@ def delete_entity_types( """ entity_types = [] for entity_type_id in entity_type_ids: - entity_type = self.get_entity_type(entity_type_id=entity_type_id) + entity_type = self._get_entity_type(entity_type_id=entity_type_id) entity_type.delete(force=force, sync=False) entity_types.append(entity_type) @@ -551,6 +563,7 @@ def create_entity_type( featurestore.EntityType - EntityType resource object """ + self.wait() return featurestore.EntityType.create( entity_type_id=entity_type_id, featurestore_name=self.resource_name, @@ -595,8 +608,9 @@ def _batch_read_feature_values( return self + @staticmethod def _validate_and_get_read_instances( - self, read_instances_uri: str, + read_instances_uri: str, ) -> Union[gca_io.BigQuerySource, gca_io.CsvSource]: """Gets read_instances @@ -629,6 +643,7 @@ def _validate_and_get_read_instances( def _validate_and_get_batch_read_feature_values_request( self, + featurestore_name: str, serving_feature_ids: Dict[str, List[str]], destination: Union[ gca_io.BigQueryDestination, @@ -642,6 +657,8 @@ def _validate_and_get_batch_read_feature_values_request( """Validates and gets batch_read_feature_values_request Args: + featurestore_name (str): + Required. A fully-qualified featurestore resource name. serving_feature_ids (Dict[str, List[str]]): Required. A user defined dictionary to define the entity_types and their features for batch serve/read. The keys of the dictionary are the serving entity_type ids and @@ -680,9 +697,7 @@ def _validate_and_get_batch_read_feature_values_request( Returns: gca_featurestore_service.BatchReadFeatureValuesRequest: batch read feature values request """ - - self.wait() - featurestore_name_components = self._parse_resource_name(self.resource_name) + featurestore_name_components = self._parse_resource_name(featurestore_name) feature_destination_fields = feature_destination_fields or {} @@ -720,7 +735,7 @@ def _validate_and_get_batch_read_feature_values_request( entity_type_specs.append(entity_type_spec) batch_read_feature_values_request = gca_featurestore_service.BatchReadFeatureValuesRequest( - featurestore=self.resource_name, entity_type_specs=entity_type_specs, + featurestore=featurestore_name, entity_type_specs=entity_type_specs, ) if isinstance(destination, gca_io.BigQueryDestination): @@ -843,6 +858,7 @@ def batch_serve_to_bq( read_instances = self._validate_and_get_read_instances(read_instances_uri) batch_read_feature_values_request = self._validate_and_get_batch_read_feature_values_request( + featurestore_name=self.resource_name, serving_feature_ids=serving_feature_ids, destination=gca_io.BigQueryDestination( output_uri=bq_destination_output_uri @@ -989,6 +1005,7 @@ def batch_serve_to_gcs( read_instances = self._validate_and_get_read_instances(read_instances_uri) batch_read_feature_values_request = self._validate_and_get_batch_read_feature_values_request( + featurestore_name=self.resource_name, serving_feature_ids=serving_feature_ids, destination=destination, feature_destination_fields=feature_destination_fields, diff --git a/tests/unit/aiplatform/test_featurestores.py b/tests/unit/aiplatform/test_featurestores.py index 5da494c443..df7d544d95 100644 --- a/tests/unit/aiplatform/test_featurestores.py +++ b/tests/unit/aiplatform/test_featurestores.py @@ -1097,7 +1097,7 @@ def test_validate_and_get_batch_read_feature_values_request( featurestore_name=_TEST_FEATURESTORE_NAME ) expected_batch_read_feature_values_request = gca_featurestore_service.BatchReadFeatureValuesRequest( - featurestore=my_featurestore.resource_name, + featurestore=_TEST_FEATURESTORE_NAME, destination=gca_featurestore_service.FeatureValueDestination( bigquery_destination=_TEST_BQ_DESTINATION, ), @@ -1107,6 +1107,7 @@ def test_validate_and_get_batch_read_feature_values_request( assert ( expected_batch_read_feature_values_request == my_featurestore._validate_and_get_batch_read_feature_values_request( + featurestore_name=my_featurestore.resource_name, serving_feature_ids=serving_feature_ids, destination=_TEST_BQ_DESTINATION, read_instances=_TEST_BQ_SOURCE, @@ -1560,6 +1561,7 @@ def test_validate_and_get_import_feature_values_request_with_source_fields(self) assert ( true_import_feature_values_request == my_entity_type._validate_and_get_import_feature_values_request( + entity_type_name=my_entity_type.resource_name, feature_ids=_TEST_IMPORTING_FEATURE_IDS, feature_time=_TEST_FEATURE_TIME_FIELD, data_source=_TEST_BQ_SOURCE, @@ -1586,6 +1588,7 @@ def test_validate_and_get_import_feature_values_request_without_source_fields(se assert ( true_import_feature_values_request == my_entity_type._validate_and_get_import_feature_values_request( + entity_type_name=my_entity_type.resource_name, feature_ids=_TEST_IMPORTING_FEATURE_IDS, feature_time=_TEST_FEATURE_TIME, data_source=_TEST_CSV_SOURCE,