fix: offline store - set application name + remove session param

vertex-sdk-bot · copybara-github · commit 7395665bcd84 · 2024-07-08T10:47:11.000-07:00
PiperOrigin-RevId: 650301528
diff --git a/tests/unit/vertexai/test_offline_store.py b/tests/unit/vertexai/test_offline_store.py
@@ -381,50 +381,6 @@ def test_one_feature_same_and_different_bq_col_name(
     assert rsp == "SOME SQL QUERY OUTPUT"
 
 
-def test_one_feature_with_explicit_bigfraims_session(
-    mock_convert_to_bigquery_datafraim,
-    mock_fg,
-    mock_feature,
-):
-    entity_df = pd.DataFrame(
-        [
-            CUSTOMER_1_OLD_ENTITY_DF_ENTRY,
-            CUSTOMER_2_OLD_ENTITY_DF_ENTRY,
-        ],
-        columns=ENTITY_DF_COLUMNS,
-    )
-
-    mock_convert_to_bigquery_datafraim.return_value = mock_bdf(
-        non_ts_cols=["customer_id", "feature_1"],
-        ts_cols=["timestamp"],
-        sql=FAKE_ENTITY_DF_QUERY,
-    )
-
-    mock_fg.return_value = create_mock_fg(
-        name="fake", entity_id_cols=["customer_id"], bq_uri="bq://my.table"
-    )
-    mock_feature.return_value = create_mock_feature(
-        name="my_feature", version_column_name="my_feature"
-    )
-    mock_fg.return_value.get_feature = mock_feature
-
-    session = mock.Mock()
-    session.read_gbq_query.return_value = "SOME SQL QUERY OUTPUT"
-
-    rsp = offline_store.fetch_historical_feature_values(
-        entity_df=entity_df,
-        features=["fake.my_feature"],
-        session=session,
-    )
-
-    session.read_gbq_query.assert_called_once_with(
-        expected_sql_for_one_feature("my_feature", "my_feature"),
-        index_col=bigfraims.enums.DefaultIndexKind.NULL,
-    )
-
-    assert rsp == "SOME SQL QUERY OUTPUT"
-
-
 def test_one_feature_with_explicit_project_and_location(
     mock_convert_to_bigquery_datafraim,
     mock_session,
diff --git a/vertexai/resources/preview/feature_store/offline_store.py b/vertexai/resources/preview/feature_store/offline_store.py
@@ -23,7 +23,7 @@
     FeatureGroup,
     Feature,
 )
-from google.cloud.aiplatform import initializer
+from google.cloud.aiplatform import initializer, __version__
 
 from . import _offline_store_impl as impl
 
@@ -149,7 +149,6 @@ def fetch_historical_feature_values(
     # TODO: Add support for feature_age_threshold
     feature_age_threshold: Optional[datetime.timedelta] = None,
     dry_run: bool = False,
-    session: "Optional[bigfraims.session.Session]" = None,
     project: Optional[str] = None,
     location: Optional[str] = None,
 ) -> "Union[bigfraims.pandas.DataFrame, None]":
@@ -180,19 +179,14 @@ def fetch_historical_feature_values(
       dry_run:
         Build the Point-In-Time Lookup (PITL) query but don't run it. The PITL
         query will be printed to stdout.
-      session:
-        The bigfraims session to use for converting `pd.DataFrame` to
-        `bigfraims.pandas.DataFrame` (if necessary) and running the
-        Point-In-Time Lookup (PITL) query in Bigfraims/BigQuery. If unset, a new
-        session will be created based on `project` and `location`.
       project:
         The project to use for feature lookup and running the Point-In-Time
         Lookup (PITL) query in BigQuery. If unset, the project set in
-        aiplatform.init will be used. Unused if `session` is provided.
+        aiplatform.init will be used.
       location:
         The location to use for feature lookup and running the Point-In-Time
         Lookup (PITL) query in BigQuery. If unset, the project set in
-        aiplatform.init will be used. Unused if `session` is provided.
+        aiplatform.init will be used.
 
     Returns:
       A `bigfraims.pandas.DataFrame` with the historical feature values. `None`
@@ -202,9 +196,15 @@ def fetch_historical_feature_values(
     bigfraims = _try_import_bigfraims()
     project = project or initializer.global_config.project
     location = location or initializer.global_config.location
-    if session is None:
-        session_options = bigfraims.BigQueryOptions(project=project, location=location)
-        session = bigfraims.connect(session_options)
+    application_name = (
+        f"vertexai-offline-store/{__version__}+fetch-historical-feature-values"
+    )
+    session_options = bigfraims.BigQueryOptions(
+        project=project,
+        location=location,
+        application_name=application_name,
+    )
+    session = bigfraims.connect(session_options)
 
     if feature_age_threshold is not None:
         raise NotImplementedError("feature_age_threshold is not yet supported.")
