fix: use the project id from BQ dataset instead of the default projec…

…t id

It is possible that the user specify a different project than the current project as the BQ output destination for batch prediction job.

google/cloud/aiplatform/jobs.py

@@ -817,11 +817,8 @@ def iter_outputs(
             if bq_dataset.startswith("bq://"):
                 bq_dataset = bq_dataset[5:]
 
-            # # Split project ID and BQ dataset ID
-            _, bq_dataset_id = bq_dataset.split(".", 1)
-
             row_iterator = bq_client.list_rows(
-                table=f"{bq_dataset_id}.predictions", max_results=bq_max_results
+                table=f"{bq_dataset}.predictions", max_results=bq_max_results
             )
 
             return row_iterator

tests/unit/aiplatform/test_jobs.py

@@ -55,12 +55,13 @@
 _TEST_ID = "1028944691210842416"
 _TEST_ALT_ID = "8834795523125638878"
 _TEST_DISPLAY_NAME = "my_job_1234"
+_TEST_BQ_PROJECT_ID = "projectId"
 _TEST_BQ_DATASET_ID = "bqDatasetId"
 _TEST_BQ_JOB_ID = "123459876"
 _TEST_BQ_MAX_RESULTS = 100
 _TEST_GCS_BUCKET_NAME = "my-bucket"
 
-_TEST_BQ_PATH = f"bq://projectId.{_TEST_BQ_DATASET_ID}"
+_TEST_BQ_PATH = f"bq://{_TEST_BQ_PROJECT_ID}.{_TEST_BQ_DATASET_ID}"
 _TEST_GCS_BUCKET_PATH = f"gs://{_TEST_GCS_BUCKET_NAME}"
 _TEST_GCS_JSONL_SOURCE_URI = f"{_TEST_GCS_BUCKET_PATH}/bp_input_config.jsonl"
 _TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}"
@@ -397,7 +398,8 @@ def test_batch_prediction_iter_dirs_bq(self, bq_list_rows_mock):
         bp.iter_outputs()
 
         bq_list_rows_mock.assert_called_once_with(
-            table=f"{_TEST_BQ_DATASET_ID}.predictions", max_results=_TEST_BQ_MAX_RESULTS
+            table=f"{_TEST_BQ_PROJECT_ID}.{_TEST_BQ_DATASET_ID}.predictions",
+            max_results=_TEST_BQ_MAX_RESULTS,
         )
 
     @pytest.mark.usefixtures("get_batch_prediction_job_running_bq_output_mock")