chore: offline store - pass credentials to bigframes session

PiperOrigin-RevId: 650355180

Author: vertex-sdk-bot

tests/unit/vertexai/test_offline_store.py

@@ -427,3 +427,55 @@ def test_one_feature_with_explicit_project_and_location(
     )
 
     assert rsp == "SOME SQL QUERY OUTPUT"
+
+
+def test_one_feature_with_explicit_credentials(
+    mock_convert_to_bigquery_dataframe,
+    mock_session,
+    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_dataframe.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
+
+    mock_session.return_value.read_gbq_query.return_value = "SOME SQL QUERY OUTPUT"
+
+    credentials = mock.MagicMock()
+    rsp = offline_store.fetch_historical_feature_values(
+        entity_df=entity_df,
+        features=["fake.my_feature"],
+        credentials=credentials,
+    )
+
+    # Ensure when initializing the session, the credentials are
+    # passed through.
+    mock_session.assert_called_once_with(
+        bigframes.BigQueryOptions(
+            credentials=credentials, project="my-project", location="my-location"
+        ),
+    )
+    mock_session.return_value.read_gbq_query.assert_called_once_with(
+        expected_sql_for_one_feature("my_feature", "my_feature"),
+        index_col=bigframes.enums.DefaultIndexKind.NULL,
+    )
+
+    assert rsp == "SOME SQL QUERY OUTPUT"

vertexai/resources/preview/feature_store/offline_store.py

@@ -19,6 +19,7 @@
 import re
 
 from typing import Optional, List, Tuple, Union, TYPE_CHECKING
+from google.auth import credentials as auth_credentials
 from vertexai.resources.preview.feature_store import (
     FeatureGroup,
     Feature,
@@ -151,6 +152,7 @@ def fetch_historical_feature_values(
     dry_run: bool = False,
     project: Optional[str] = None,
     location: Optional[str] = None,
+    credentials: Optional[auth_credentials.Credentials] = None,
 ) -> "Union[bigframes.pandas.DataFrame, None]":
     """Fetch historical data at the timestamp specified for each entity.
 
@@ -187,6 +189,9 @@ def fetch_historical_feature_values(
         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.
+      credentials:
+        Custom credentials to use for running the Point-In-Time Lookup (PITL)
+        query in BigQuery. Overrides credentials set in aiplatform.init.
 
     Returns:
       A `bigframes.pandas.DataFrame` with the historical feature values. `None`
@@ -196,10 +201,12 @@ def fetch_historical_feature_values(
     bigframes = _try_import_bigframes()
     project = project or initializer.global_config.project
     location = location or initializer.global_config.location
+    credentials = credentials or initializer.global_config.credentials
     application_name = (
         f"vertexai-offline-store/{__version__}+fetch-historical-feature-values"
     )
     session_options = bigframes.BigQueryOptions(
+        credentials=credentials,
         project=project,
         location=location,
         application_name=application_name,