From 0664ea32eff7be779968e082a8b7e85f72c4791f Mon Sep 17 00:00:00 2001 From: Frances Hubis Thoma Date: Mon, 7 Jul 2025 11:09:44 -0700 Subject: [PATCH] feat: Configure Bigframes implicitly in `MultimodalDataset.assess()`. PiperOrigin-RevId: 780182958 --- google/cloud/aiplatform/preview/datasets.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/google/cloud/aiplatform/preview/datasets.py b/google/cloud/aiplatform/preview/datasets.py index 485645b31d..45eaa1e0b0 100644 --- a/google/cloud/aiplatform/preview/datasets.py +++ b/google/cloud/aiplatform/preview/datasets.py @@ -1402,10 +1402,18 @@ def assemble( result = assemble_lro.result(timeout=None) _LOGGER.log_action_completed_against_resource("data", "assembled", self) table_id = result.bigquery_destination.lstrip("bq://") - return ( - table_id, - bigframes.pandas.read_gbq(table_id) if load_dataframe else None, - ) + if load_dataframe: + session_options = bigframes.BigQueryOptions( + credentials=initializer.global_config.credentials, + project=initializer.global_config.project, + location=initializer.global_config.location, + ) + with bigframes.connect(session_options) as session: + df = session.read_gbq(table_id) + else: + df = None + + return (table_id, df) def assess_tuning_resources( self,